diff --git a/simpeg/dask/electromagnetics/time_domain/simulation.py b/simpeg/dask/electromagnetics/time_domain/simulation.py index 5986f57ffa..9accf6d22b 100644 --- a/simpeg/dask/electromagnetics/time_domain/simulation.py +++ b/simpeg/dask/electromagnetics/time_domain/simulation.py @@ -88,10 +88,8 @@ def compute_J(self, m, f=None): f, Ainv = self.fields(m=m, return_Ainv=True) client, worker = self._get_client_worker() - ftype = self._fieldType + "Solution" n_cells = m.size - simulation_times = np.r_[0, np.cumsum(self.time_steps)] + self.t0 data_times = self.survey.source_list[0].receiver_list[0].times compute_row_size = np.ceil(self.max_chunk_size / (m.shape[0] * 8.0 * 1e-6)) @@ -99,6 +97,7 @@ def compute_J(self, m, f=None): self.survey.source_list, compute_row_size, thread_count=self.n_threads(client=client, worker=worker), + optimize=True, ) fields_array = f[:, ftype, :] @@ -156,7 +155,6 @@ def compute_J(self, m, f=None): AdiagTinv, ATinv_df_duT_v[ind], time_mask, - client, ) if client: @@ -166,7 +164,7 @@ def compute_J(self, m, f=None): for block_ind in range(len(blocks)): - if len(block) == 0: + if len(blocks[block_ind]) == 0: continue if client: @@ -337,48 +335,53 @@ def get_field_deriv_block( AdiagTinv, ATinv_df_duT_v, time_mask, - client, ): """ Stack the blocks of field derivatives for a given timestep and call the direct solver. """ - if len(ATinv_df_duT_v) == 0: - ATinv_df_duT_v = [[] for _ in block] + if len(block) == 0: + return None Asubdiag = None if tInd < self.nT - 1: Asubdiag = self.getAsubdiag(tInd + 1) - updated_ATinv_df_duT_v = [] - - for (_, (rx_ind, _, shape)), field_deriv, ATinv_chunk in zip( - block, field_derivs, ATinv_df_duT_v - ): + time_blocks = [] + colm_indices = [] + colm_count = 0 + for (_, (rx_ind, _, shape)), field_deriv in zip(block, field_derivs): # Cut out early data time_check = np.kron(time_mask, np.ones(shape, dtype=bool))[rx_ind] local_ind = np.arange(rx_ind.shape[0])[time_check] - if len(ATinv_chunk) == 0: + if len(ATinv_df_duT_v) == 0: # last timestep (first to be solved) time_block = field_deriv.toarray()[:, local_ind] - shape = ( - field_deriv.shape[0], - len(rx_ind), - ) - ATinv_chunk = np.zeros(shape, dtype=np.float32) else: time_block = np.asarray( - field_deriv[:, local_ind] - Asubdiag.T * ATinv_chunk[:, local_ind] + field_deriv[:, local_ind] + - Asubdiag.T + * ATinv_df_duT_v[:, colm_count : colm_count + rx_ind.shape[0]][ + :, local_ind + ] ) - if time_block.ndim == 2 and time_block.shape[1] > 0: - solve = (AdiagTinv * time_block).reshape(time_block.shape) - ATinv_chunk[:, local_ind] = solve + time_blocks.append(time_block) + colm_indices.append(local_ind + colm_count) + colm_count += rx_ind.shape[0] + + if len(ATinv_df_duT_v) == 0: + ATinv_df_duT_v = np.zeros((field_deriv.shape[0], colm_count), dtype=np.float32) - updated_ATinv_df_duT_v.append(ATinv_chunk) + block_array = np.hstack(time_blocks).reshape((ATinv_df_duT_v.shape[0], -1)) + if block_array.shape[1] > 0: + solve = AdiagTinv * block_array + ATinv_df_duT_v[:, np.hstack(colm_indices).flatten()] = solve.reshape( + (-1, block_array.shape[1]) + ) - return updated_ATinv_df_duT_v + return ATinv_df_duT_v def block_deriv( @@ -462,11 +465,14 @@ def compute_rows( Compute the rows of the sensitivity matrix for a given source and receiver. """ rows = [] - for ind, (address, ind_array) in enumerate(blocks[block_ind]): + colm_count = 0 + for address, ind_array in blocks[block_ind]: # for (address, ind_array), field_derivs in zip(chunks, ATinv_df_duT_v): src = simulation.survey.source_list[address[0]] time_check = np.kron(time_mask, np.ones(ind_array[2], dtype=bool))[ind_array[0]] - local_ind = np.arange(len(ind_array[0]))[time_check] + + n_rec = len(ind_array[0]) + local_ind = np.arange(n_rec)[time_check] if len(local_ind) < 1: row_block = np.zeros( @@ -478,18 +484,24 @@ def compute_rows( dAsubdiagT_dm_v = simulation.getAsubdiagDeriv( tInd, fields[:, address[0], tInd], - field_derivs[block_ind][ind][:, local_ind], + field_derivs[block_ind][:, colm_count : colm_count + n_rec][:, local_ind], adjoint=True, ) dRHST_dm_v = simulation.getRHSDeriv( - tInd + 1, src, field_derivs[block_ind][ind][:, local_ind], adjoint=True + tInd + 1, + src, + field_derivs[block_ind][:, colm_count : colm_count + n_rec][:, local_ind], + adjoint=True, ) # on nodes of time mesh un_src = fields[:, address[0], tInd + 1] # cell centered on time mesh dAT_dm_v = simulation.getAdiagDeriv( - tInd, un_src, field_derivs[block_ind][ind][:, local_ind], adjoint=True + tInd, + un_src, + field_derivs[block_ind][:, colm_count : colm_count + n_rec][:, local_ind], + adjoint=True, ) row_block = np.zeros( (len(ind_array[1]), simulation.model.size), dtype=np.float32 @@ -506,6 +518,8 @@ def compute_rows( else: Jmatrix[ind_array[1], :] += row_block + colm_count += n_rec + def evaluate_dpred_block(indices, sources, mesh, time_mesh, fields): """ diff --git a/simpeg/dask/potential_fields/base.py b/simpeg/dask/potential_fields/base.py index 20a3ee8dc6..d1589abb1d 100644 --- a/simpeg/dask/potential_fields/base.py +++ b/simpeg/dask/potential_fields/base.py @@ -4,10 +4,8 @@ import os from dask import delayed, array, compute - from dask.diagnostics import ProgressBar - -import zarr +from dask.distributed import Client _chunk_format = "row" @@ -41,7 +39,7 @@ def residual(self, m, dobs, f=None): return self.dpred(m, f=f) - dobs -def block_compute(sim, rows, components, j_matrix, count): +def block_compute(sim, rows, components): block = [] for row in rows: block.append(sim.evaluate_integral(row, components)) @@ -49,34 +47,7 @@ def block_compute(sim, rows, components, j_matrix, count): if sim.store_sensitivities == "forward_only": return np.hstack(block) - values = np.vstack(block) - return storage_formatter(values, count, j_matrix) - - -def storage_formatter( - rows: np.ndarray, - count: int, - j_matrix: zarr.Array | None = None, -): - """ - Format the storage of the sensitivity matrix. - - :param rows: List of dask arrays representing blocks of the sensitivity matrix. - :param count: Current row count offset. - :param j_matrix: Zarr array to store the sensitivity matrix on disk, if applicable - - :return: If j_matrix is provided, returns None after storing the rows; otherwise, - returns the stacked rows as a NumPy array. - """ - - if isinstance(j_matrix, zarr.Array): - j_matrix.set_orthogonal_selection( - (np.arange(count, count + rows.shape[0]), slice(None)), - rows.astype(np.float32), - ) - return None - - return rows + return np.vstack(block) def linear_operator(self): @@ -86,44 +57,34 @@ def linear_operator(self): n_cells *= 3 if self.store_sensitivities == "disk": - if os.path.exists(self.sensitivity_path): return array.from_zarr(self.sensitivity_path) - Jmatrix = zarr.open( - self.sensitivity_path, - mode="w", - shape=(self.survey.nD, n_cells), - chunks=(self.max_chunk_size, n_cells), - ) - else: - Jmatrix = None - n_components = len(self.survey.source_list[0].receiver_list[0].components) n_blocks = np.ceil( (n_cells * n_components * self.survey.receiver_locations.shape[0] * 8.0 * 1e-6) / self.max_chunk_size ) block_split = np.array_split(self.survey.receiver_locations, n_blocks) - client, worker = self._get_client_worker() - if client: + if client is None: + client = Client() + + if client and worker: sim = client.scatter(self, workers=worker) else: delayed_compute = delayed(block_compute) rows = [] count = 0 - for block in block_split: - if client: + for count, block in enumerate(block_split): + if client and worker: row = client.submit( block_compute, sim, block, self.survey.source_list[0].receiver_list[0].components, - Jmatrix, - count, workers=worker, ) @@ -132,8 +93,6 @@ def linear_operator(self): self, block, self.survey.source_list[0].receiver_list[0].components, - Jmatrix, - count, ) row = array.from_delayed( chunk, @@ -147,14 +106,22 @@ def linear_operator(self): count += block.shape[0] rows.append(row) - if client: + if client and worker: kernel = client.gather(rows) - else: + elif forward_only: with ProgressBar(): kernel = compute(rows)[0] + else: + kernel = rows - if self.store_sensitivities == "disk" and os.path.exists(self.sensitivity_path): - return array.from_zarr(self.sensitivity_path) + if self.store_sensitivities == "disk": + j_matrix = array.concatenate(rows, axis=0) + + with ProgressBar(): + j_matrix = j_matrix.to_zarr( + self.sensitivity_path, return_stored=True, compute=True + ) + return j_matrix if forward_only: return np.hstack(kernel)