View Javadoc
1   // ******************************************************************************
2   //
3   // Title:       Force Field X.
4   // Description: Force Field X - Software for Molecular Biophysics.
5   // Copyright:   Copyright (c) Michael J. Schnieders 2001-2025.
6   //
7   // This file is part of Force Field X.
8   //
9   // Force Field X is free software; you can redistribute it and/or modify it
10  // under the terms of the GNU General Public License version 3 as published by
11  // the Free Software Foundation.
12  //
13  // Force Field X is distributed in the hope that it will be useful, but WITHOUT
14  // ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
15  // FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
16  // details.
17  //
18  // You should have received a copy of the GNU General Public License along with
19  // Force Field X; if not, write to the Free Software Foundation, Inc., 59 Temple
20  // Place, Suite 330, Boston, MA 02111-1307 USA
21  //
22  // Linking this library statically or dynamically with other modules is making a
23  // combined work based on this library. Thus, the terms and conditions of the
24  // GNU General Public License cover the whole combination.
25  //
26  // As a special exception, the copyright holders of this library give you
27  // permission to link this library with independent modules to produce an
28  // executable, regardless of the license terms of these independent modules, and
29  // to copy and distribute the resulting executable under terms of your choice,
30  // provided that you also meet, for each linked independent module, the terms
31  // and conditions of the license of that module. An independent module is a
32  // module which is not derived from or based on this library. If you modify this
33  // library, you may extend this exception to your version of the library, but
34  // you are not obligated to do so. If you do not wish to do so, delete this
35  // exception statement from your version.
36  //
37  // ******************************************************************************
38  package ffx.numerics.fft;
39  
40  /**
41   * Compute the 2D FFT of complex, double precision input of arbitrary dimensions
42   * via 1D Mixed Radix FFTs.
43   * <p>
44   * For interleaved data, the location of the input point [x, y] within the input array must be: <br>
45   * double real = input[x*nextX + y*nextY]<br>
46   * double imag = input[x*nextX + y*nextY + im]<br>
47   * where <br>
48   * nextX = 2
49   * nextY = 2*nX
50   * im = 1
51   * <p>
52   * For blocked data along x, the location of the input point [x, y] within the input array must be: <br>
53   * double real = input[x*nextX + y*nextY]<br>
54   * double imag = input[x*nextX + y*nextY + im]<br>
55   * where for BLOCKED_X <br>
56   * nextX = 1<br>
57   * nextY = nX<br>
58   * im = nX<br>
59   * and for BLOCKED_XY <br>
60   * nextX = 1<br>
61   * nextY = nX*nY<br>
62   * im = nX*nY<br>
63   *
64   * @author Michal J. Schnieders
65   * @see Complex
66   * @since 1.0
67   */
68  public class Complex2D {
69  
70    /**
71     * The 2D data layout.
72     */
73    private final DataLayout2D layout;
74    /**
75     * The external imaginary offset.
76     */
77    private final int externalIm;
78    /**
79     * The X-dimension.
80     */
81    private final int nX;
82    /**
83     * The Y-dimension.
84     */
85    private final int nY;
86    /**
87     * The next real value along the X dimension.
88     */
89    private final int nextX;
90    /**
91     * The next real value along the Y dimension.
92     */
93    private final int nextY;
94    /**
95     * Compute FFTs along X one at a time.
96     */
97    private final Complex fftX;
98    /**
99     * Compute FFTs along Y one at a time.
100    */
101   private final Complex fftY;
102   /**
103    * If true, pack FFTs along X (or Y) into a contiguous array to compute all FFTs along X (or Y) at once.
104    */
105   private boolean packFFTs;
106   /**
107    * If true, use SIMD instructions.
108    */
109   private boolean useSIMD;
110   /**
111    * Compute nY FFTs along the X dimension all at once.
112    */
113   private final Complex packedFFTX;
114   /**
115    * Compute nX FFTs along the Y dimension all at once.
116    */
117   private final Complex packedFFTY;
118   /**
119    * Working array for packed FFTs.
120    */
121   private final double[] packedData;
122   /**
123    * The offset between real values in the packed data.
124    */
125   private final int ii;
126   /**
127    * The offset between any real value and its corresponding imaginary value for the packed data.
128    */
129   private final int im;
130   /**
131    * The offset between real values along the X-dimension in the transposed packed data.
132    */
133   private final int trNextX;
134   /**
135    * The offset between real values along the Y-dimension in the transposed packed data.
136    */
137   private final int trNextY;
138 
139   /**
140    * Create a new 2D Complex FFT for interleaved data.
141    *
142    * @param nX The number of points in the X dimension.
143    * @param nY The number of points in the Y dimension.
144    */
145   public Complex2D(int nX, int nY) {
146     this(nX, nY, DataLayout2D.INTERLEAVED, 1);
147   }
148 
149   /**
150    * Create a new 2D Complex FFT.
151    *
152    * @param nX       The number of points in the X dimension.
153    * @param nY       The number of points in the Y dimension.
154    * @param layout   The data layout.
155    * @param imOffset The offset between real and imaginary values.
156    */
157   public Complex2D(int nX, int nY, DataLayout2D layout, int imOffset) {
158     this.nX = nX;
159     this.nY = nY;
160     this.externalIm = imOffset;
161     this.layout = layout;
162     if (layout == DataLayout2D.INTERLEAVED) {
163       im = 1;
164       ii = 2;
165       nextX = 2;
166       nextY = 2 * nX;
167       trNextY = 2;
168       trNextX = 2 * nY;
169     } else if (layout == DataLayout2D.BLOCKED_X) {
170       im = nX;
171       ii = 1;
172       nextX = 1;
173       nextY = 2 * nX;
174       trNextY = 1;
175       trNextX = 2 * nY;
176     } else if (layout == DataLayout2D.BLOCKED_XY) {
177       im = nX * nY;
178       ii = 1;
179       nextX = 1;
180       nextY = nX;
181       trNextY = 1;
182       trNextX = nY;
183     } else {
184       throw new IllegalArgumentException("Unsupported data layout: " + layout);
185     }
186 
187     if (this.externalIm != im) {
188       throw new IllegalArgumentException("Unsupported im offset: " + imOffset);
189     }
190 
191     // Do not use SIMD by default for now.
192     useSIMD = false;
193     String simd = System.getProperty("fft.useSIMD", Boolean.toString(useSIMD));
194     try {
195       useSIMD = Boolean.parseBoolean(simd);
196     } catch (Exception e) {
197       useSIMD = false;
198     }
199 
200     packFFTs = false;
201     String pack = System.getProperty("fft.packFFTs", Boolean.toString(packFFTs));
202     try {
203       packFFTs = Boolean.parseBoolean(pack);
204     } catch (Exception e) {
205       packFFTs = false;
206     }
207 
208     DataLayout1D layout1D;
209     if (layout == DataLayout2D.INTERLEAVED) {
210       layout1D = DataLayout1D.INTERLEAVED;
211     } else {
212       layout1D = DataLayout1D.BLOCKED;
213     }
214 
215     // Create 1D FFTs that will be used for computing 1 FFT at a time.
216     fftX = new Complex(nX, layout1D, externalIm);
217     fftX.setUseSIMD(useSIMD);
218     fftY = new Complex(nY, layout1D, externalIm);
219     fftY.setUseSIMD(useSIMD);
220 
221     // Create 1D FFTs that will be used for computing all FFTs at once.
222     packedFFTY = new Complex(nY, layout1D, externalIm, nX);
223     packedFFTY.setUseSIMD(useSIMD);
224     packedFFTX = new Complex(nX, layout1D, im, nY);
225     packedFFTX.setUseSIMD(useSIMD);
226     packedData = new double[2 * nX * nY];
227   }
228 
229   /**
230    * Get the Data Layout.
231    *
232    * @return The layout.
233    */
234   public DataLayout2D getLayout() {
235     return layout;
236   }
237 
238   /**
239    * Set the 2D transform to use SIMD instructions.
240    *
241    * @param useSIMD True to use SIMD instructions.
242    */
243   public void setUseSIMD(boolean useSIMD) {
244     this.useSIMD = useSIMD;
245     fftX.setUseSIMD(useSIMD);
246     fftY.setUseSIMD(useSIMD);
247     packedFFTX.setUseSIMD(useSIMD);
248     packedFFTY.setUseSIMD(useSIMD);
249   }
250 
251   /**
252    * Set the 2D transform to pack FFTs into a contiguous array to compute all FFTs at once.
253    *
254    * @param packFFTs True to pack FFTs.
255    */
256   public void setPackFFTs(boolean packFFTs) {
257     this.packFFTs = packFFTs;
258   }
259 
260   /**
261    * Compute the 2D FFT.
262    *
263    * @param input The input array must be of size 2 * nX * nY.
264    * @param index The offset into the input array of the first element.
265    */
266   public void fft(final double[] input, int index) {
267     if (!packFFTs) {
268       for (int offset = index, y = 0; y < nY; y++, offset += nextY) {
269         fftX.fft(input, offset, nextX);
270       }
271       for (int offset = index, x = 0; x < nX; x++, offset += nextX) {
272         fftY.fft(input, offset, nextY);
273       }
274     } else {
275       // For FFTs along Y, the input data is already contiguous.
276       packedFFTY.fft(input, index, ii);
277       transpose(input, index);
278       packedFFTX.fft(packedData, 0, ii);
279       untranspose(input, index);
280     }
281   }
282 
283   /**
284    * Pack the input array for Fourier transforms along the X dimension.
285    * <p>
286    * Input order:
287    * real(x,y) = input[offset + x*nextX + y*nextY]
288    * imag(x,y) = input[offset + x*nextX + y*nextY + im]
289    * Output order:
290    * real(x,y) = packedData[y*trNextY + x*trNextX]
291    * imag(x,y) = packedData[y*trNextY + x*trXextX + im]
292    *
293    * @param input  The input data.
294    * @param offset The offset into the input data.
295    */
296   private void transpose(final double[] input, int offset) {
297     int index = 0;
298     // Outer loop over the X dimension.
299     for (int x = 0; x < nX; x++) {
300       int dx = offset + x * nextX;
301       // Inner loop over the Y dimension (the number of FFTs).
302       for (int y = 0; y < nY; y++) {
303         double real = input[dx + y * nextY];
304         double imag = input[dx + y * nextY + externalIm];
305         // Contiguous storage into the packed array.
306         packedData[index] = real;
307         packedData[index + im] = imag;
308         index += ii;
309       }
310     }
311   }
312 
313   /**
314    * Unpack the output array after Fourier transforms.
315    * <p>
316    * Input order:
317    * real_xy = packedData[y*trNextY + x*trNextX]
318    * imag_xy = packedData[y*trNextY + x*trXextX + im]
319    * Output order:
320    * real_xy = output[offset + x*nextX + y*nextY]
321    * imag_xy = output[offset + x*nextX + y*nextY + im]
322    *
323    * @param output The output data.
324    * @param offset The offset into the output data.
325    */
326   private void untranspose(final double[] output, int offset) {
327     int index = offset;
328     // Outer loop over the Y dimension.
329     for (int y = 0; y < nY; y++) {
330       int dy = y * trNextY;
331       // Inner loop over the X dimension.
332       for (int x = 0; x < nX; x++) {
333         double real = packedData[dy + x * trNextX];
334         double imag = packedData[dy + x * trNextX + im];
335         // Contiguous storage into the output array.
336         output[index] = real;
337         output[index + externalIm] = imag;
338         index += ii;
339       }
340     }
341   }
342 
343   /**
344    * Compute the 2D IFFT.
345    *
346    * @param input The input array must be of size 2 * nX * nY.
347    * @param index The offset into the input array of the first element.
348    */
349   public void ifft(final double[] input, int index) {
350     if (!packFFTs) {
351       for (int offset = index, y = 0; y < nY; y++, offset += nextY) {
352         fftX.ifft(input, offset, nextX);
353       }
354       for (int offset = index, x = 0; x < nX; x++, offset += nextX) {
355         fftY.ifft(input, offset, nextY);
356       }
357     } else {
358       // For FFTs along Y, the input data is already contiguous.
359       packedFFTY.ifft(input, index, ii);
360       transpose(input, index);
361       packedFFTX.ifft(packedData, 0, ii);
362       untranspose(input, index);
363     }
364   }
365 }