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 }