| 58 | self.scaleby *= self.phase |
| 59 | |
| 60 | def handle_dataset(self, in_arr): |
| 61 | |
| 62 | out_dims = [1, 1, 1] |
| 63 | rank = len(in_arr.shape) |
| 64 | num_ones = 3 - rank |
| 65 | in_dims = list(in_arr.shape) + [1] * num_ones |
| 66 | |
| 67 | if np.iscomplexobj(in_arr): |
| 68 | in_arr_re = np.real(in_arr) |
| 69 | in_arr_im = np.imag(in_arr) |
| 70 | else: |
| 71 | in_arr_re = in_arr |
| 72 | in_arr_im = None |
| 73 | |
| 74 | if self.verbose: |
| 75 | fmt = "Input data is rank {}, size {}x{}x{}." |
| 76 | print(fmt.format(rank, in_dims[0], in_dims[1], in_dims[2])) |
| 77 | |
| 78 | if self.resolution > 0: |
| 79 | out_dims[0] = math.floor(self.Rout.c1.norm() * self.resolution + 0.5) |
| 80 | out_dims[1] = math.floor(self.Rout.c2.norm() * self.resolution + 0.5) |
| 81 | out_dims[2] = math.floor(self.Rout.c3.norm() * self.resolution + 0.5) |
| 82 | else: |
| 83 | for i in range(3): |
| 84 | out_dims[i] = in_dims[i] * self.multiply_size[i] |
| 85 | |
| 86 | for i in range(rank, 3): |
| 87 | out_dims[i] = 1 |
| 88 | |
| 89 | N = 1 |
| 90 | for i in range(3): |
| 91 | out_dims[i] = int(max(out_dims[i], 1)) |
| 92 | N *= out_dims[i] |
| 93 | |
| 94 | if self.verbose: |
| 95 | print(f"Output data {out_dims[0]}x{out_dims[1]}x{out_dims[2]}") |
| 96 | |
| 97 | out_arr_re = np.zeros(int(N)) |
| 98 | |
| 99 | if isinstance(in_arr_im, np.ndarray): |
| 100 | out_arr_im = np.zeros(int(N)) |
| 101 | else: |
| 102 | out_arr_im = np.array([]) |
| 103 | |
| 104 | flat_in_arr_re = in_arr_re.ravel() |
| 105 | flat_in_arr_im = ( |
| 106 | in_arr_im.ravel() if isinstance(in_arr_im, np.ndarray) else np.array([]) |
| 107 | ) |
| 108 | |
| 109 | kvector = [self.kpoint.x, self.kpoint.y, self.kpoint.z] if self.kpoint else [] |
| 110 | map_data( |
| 111 | flat_in_arr_re, |
| 112 | flat_in_arr_im, |
| 113 | np.array(in_dims, dtype=np.intc), |
| 114 | out_arr_re, |
| 115 | out_arr_im, |
| 116 | np.array(out_dims, dtype=np.intc), |
| 117 | self.coord_map, |