diff --git a/simpeg/dask/electromagnetics/time_domain/simulation.py b/simpeg/dask/electromagnetics/time_domain/simulation.py index 5986f57ffa..07b12a2c61 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( @@ -461,35 +464,35 @@ 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] - - if len(local_ind) < 1: - row_block = np.zeros( - (len(ind_array[1]), simulation.model.size), dtype=np.float32 - ) - rows.append(row_block) - continue + n_rec = len(ind_array[0]) + local_ind = np.arange(n_rec)[time_check] 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 +509,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 05f4231141..23df261b31 100644 --- a/simpeg/dask/potential_fields/base.py +++ b/simpeg/dask/potential_fields/base.py @@ -5,8 +5,6 @@ import os from dask import delayed, array, compute from dask.diagnostics import ProgressBar -from dask.distributed import Client - _chunk_format = "row" @@ -68,9 +66,7 @@ def linear_operator(self): block_split = np.array_split(self.survey.receiver_locations, n_blocks) client, worker = self._get_client_worker() - if client is None: - client = Client() - + # If no distributed Client is active, fall back to the local scheduler. if client and worker and self.store_sensitivities != "disk": sim = client.scatter(self, workers=worker) else: @@ -110,20 +106,19 @@ def linear_operator(self): if client and worker: kernel = client.gather(rows) - elif forward_only: + elif self.store_sensitivities == "disk": + kernel = rows + else: with ProgressBar(): kernel = compute(rows)[0] - else: - kernel = rows 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 + j_matrix.to_zarr(self.sensitivity_path, compute=True) + + return array.from_zarr(self.sensitivity_path) if forward_only: return np.hstack(kernel)