diff --git a/tests/warmth3d/test_3d.py b/tests/warmth3d/test_3d.py index 61db873..c7a3c62 100644 --- a/tests/warmth3d/test_3d.py +++ b/tests/warmth3d/test_3d.py @@ -11,7 +11,7 @@ @pytest.mark.mpi def test_3d_compare(): comm = MPI.COMM_WORLD - inc = 2100 + inc = 1000 model_pickled = f"model-out-inc_{inc}.p" if comm.rank == 0 and not os.path.isfile(model_pickled): global runtime_1D_sim @@ -46,7 +46,7 @@ def test_3d_compare(): i.adiab = 0.3e-3 - model.simulator.simulate_every = 1 + model.simulator.simulate_every = 2 # # set 1D simulation parameters to be most similar to those in the (later) 3D simulation, for better comparison @@ -62,7 +62,6 @@ def test_3d_compare(): print("Total time 1D simulations:", runtime_1D_sim) pickle.dump( model, open( model_pickled, "wb" ) ) - # model = pickle.load( open( model_pickled, "rb" ) ) try: os.mkdir('mesh') except FileExistsError: @@ -84,7 +83,12 @@ def test_3d_compare(): hx = nnx // 2 hy = nny // 2 - nn = model.builder.nodes[hy-mm2.padX][hx-mm2.padX] + nn0 = model.builder.nodes[hy-mm2.padX][hx-mm2.padX] + + node_result_path = str(nn0.node_path).replace(".pickle", "_results") + assert os.path.exists(node_result_path), f"ERROR: Node result file {node_result_path} is missing." + + nn = pickle.load(open(node_result_path,"rb")) dd = nn._depth_out[:,0] mm2_pos, mm2_temp = mm2.get_node_pos_and_temp(-1) diff --git a/warmth/build.py b/warmth/build.py index bc5be59..ac89b40 100644 --- a/warmth/build.py +++ b/warmth/build.py @@ -108,7 +108,10 @@ def __init__(self): self._crust_ls:np.ndarray[np.float64]|None=None self._lith_ls:np.ndarray[np.float64]|None=None self._subsidence:np.ndarray[np.float64]|None=None - + self.seabed_arr:np.ndarray[np.float64]|None=None + self.top_crust_arr:np.ndarray[np.float64]|None=None + self.top_lith_arr:np.ndarray[np.float64]|None=None + self.top_aest_arr:np.ndarray[np.float64]|None=None @property def shf(self)->float: @@ -123,18 +126,37 @@ def result(self)-> Results|None: Results|None None if not simulated """ - items = [self._depth_out,self.temperature_out,self._idsed] - if any(isinstance(i,type(None)) for i in items): - return None - else: - return Results(self._depth_out,self.temperature_out,self._idsed,self.sediments,self.kCrust,self.kLith,self.kAsth) + return Results(self._depth_out,self.temperature_out,self._idsed,self.sediments,self.kCrust,self.kLith,self.kAsth) + + def clear_unused_data(self): + """Removes most arrays of detailed input and output that are not needed by warmth3D, in order to save memory. + """ + # self.max_time = self._depth_out.shape[1] + self._depth_out = None + self.temperature_out =None + self._idsed = None + self.coord_initial = None + self._crust_ls = None + self._lith_ls = None + self._subsidence =None + + def compute_derived_arrays(self): + """Computes depths of seabed, top crust, top lithosphere and top aestenosphere, and stores them with the node. + This allows the depth and temperature arrays to be discarded to save memory. + """ + self.top_crust_arr = [ self._depth_out[ np.where(self._idsed[:,age] == -1)[0][0], age] for age in range(self.max_time)] + self.top_lith_arr = [ self._depth_out[ np.where(self._idsed[:,age] == -2)[0][0], age] for age in range(self.max_time)] + self.top_aest_arr = [ self._depth_out[ np.where(self._idsed[:,age] == -3)[0][0], age] for age in range(self.max_time)] + self.seabed_arr = np.array( [ self._depth_out[np.where(~np.isnan(self.temperature_out[:,age]))[0][0],age] for age in range(self.max_time)]) + + @property def crust_ls(self)->np.ndarray[np.float64]: if isinstance(self.result,Results): - all_age = self.result.ages + all_age = np.arange(len(self.top_lith_arr),dtype=np.int32) val = np.zeros(all_age.size) for age in all_age: - val[age] = self.result.crust_thickness(age) + val[age] = self.top_lith_arr[age] - self.top_crust_arr[age] return val else: @@ -142,31 +164,34 @@ def crust_ls(self)->np.ndarray[np.float64]: @property def lith_ls(self)->np.ndarray[np.float64]: if isinstance(self.result,Results): - all_age = self.result.ages + all_age = np.arange(len(self.top_lith_arr),dtype=np.int32) val = np.zeros(all_age.size) for age in all_age: - val[age] = self.result.lithosphere_thickness(age) + val[age] = self.top_aest_arr[age] - self.top_lith_arr[age] return val else: return self._lith_ls @property def subsidence(self)->np.ndarray[np.float64]: if isinstance(self.result,Results): - all_age = self.result.ages + all_age = np.arange(len(self.top_lith_arr),dtype=np.int32) val = np.zeros(all_age.size) for age in all_age: - val[age] = self.result.seabed(age) + val[age] = self.seabed_arr[age] + # val[age] = self.result.seabed(age) return val else: return self._subsidence @property def sed_thickness_ls(self)->float: if isinstance(self.result,Results): - all_age = self.result.ages + all_age = np.arange(len(self.top_lith_arr),dtype=np.int32) val = np.zeros(all_age.size) for age in all_age: - seabed = self.result.seabed(age) - top_crust = self.result.top_crust(age) + # seabed = self.result.seabed(age) + # top_crust = self.result.top_crust(age) + seabed = self.seabed_arr[age] + top_crust = self.top_crust_arr[age] val[age] = top_crust - seabed return val else: diff --git a/warmth/forward_modelling.py b/warmth/forward_modelling.py index 006a420..9ee62de 100644 --- a/warmth/forward_modelling.py +++ b/warmth/forward_modelling.py @@ -115,7 +115,7 @@ def _sediment_conductivity_sekiguchi(mean_porosity: np.ndarray[np.float64], cond temperature_K=273.15+mid_pt_temperautureC conductivity = 1.84+358*((1.0227*conductivity)-1.882)*((1/temperature_K)-0.00068) effective_conductivity = conductivity*(1-mean_porosity) - return effective_conductivity-0.5 + return effective_conductivity def _check_beta(self, wd_diff: float, beta_current: float, beta_all: np.ndarray[np.float64], Wd_diff_all: np.ndarray[np.float64]) -> tuple[bool, np.ndarray[np.float64], np.ndarray[np.float64]]: """Check if current beta factor matches the observed subsidence diff --git a/warmth/mesh_utils.py b/warmth/mesh_utils.py index c15f34c..8a6de6b 100644 --- a/warmth/mesh_utils.py +++ b/warmth/mesh_utils.py @@ -138,13 +138,21 @@ def interpolateNode(interpolationNodes: List[single_node], interpolationWeights= node.X = np.sum( np.array( [node.X * w for node,w in zip(interpolationNodes,iWeightNorm)] ) ) node.Y = np.sum( np.array( [node.Y * w for node,w in zip(interpolationNodes,iWeightNorm)] ) ) - times = range(node.result._depth.shape[1]) + # self.top_crust_arr = [ self._depth_out[ np.where(self._idsed[:,age] == -1)[0][0], age] for age in range(self.max_time)] + # #print ("PING B") + # self.top_lith_arr = [ self._depth_out[ np.where(self._idsed[:,age] == -2)[0][0], age] for age in range(self.max_time)] + # #print ("PING C") + # self.top_aest_arr = [ self._depth_out[ np.where(self._idsed[:,age] == -3)[0][0], age] for age in range(self.max_time)] + # #print ("PING D") + + # self.top_lithosphere(age)-self.top_crust(age) + if node.subsidence is None: - node.subsidence = np.sum( np.array( [ [node.result.seabed(t) for t in times] * w for node,w in zip(interpolationNodes,iWeightNorm)] ) , axis = 0) + node.subsidence = np.sum( np.array( [ node.seabed_arr[:] * w for node,w in zip(interpolationNodes,iWeightNorm)] ) , axis = 0) if node.crust_ls is None: - node.crust_ls = np.sum( np.array( [ [node.result.crust_thickness(t) for t in times] * w for node,w in zip(interpolationNodes,iWeightNorm)] ) , axis = 0) + node.crust_ls = np.sum( np.array( [ (node.top_lith_arr[:]-node.top_crust_arr[:]) * w for node,w in zip(interpolationNodes,iWeightNorm)] ) , axis = 0) if node.lith_ls is None: - node.lith_ls = np.sum( np.array( [ [node.result.lithosphere_thickness(t) for t in times] * w for node,w in zip(interpolationNodes,iWeightNorm)] ) , axis = 0) + node.crust_ls = np.sum( np.array( [ (node.top_aest_arr[:]-node.top_lithosphere[:]) * w for node,w in zip(interpolationNodes,iWeightNorm)] ) , axis = 0) if node.beta is None: node.beta = np.sum( np.array( [node.beta * w for node,w in zip(interpolationNodes,iWeightNorm)] ) , axis = 0) @@ -152,10 +160,10 @@ def interpolateNode(interpolationNodes: List[single_node], interpolationWeights= node.kAsth = np.sum( np.array( [node.kAsth * w for node,w in zip(interpolationNodes,iWeightNorm)] ) , axis = 0) if node.kLith is None: node.kLith = np.sum( np.array( [node.kLith * w for node,w in zip(interpolationNodes,iWeightNorm)] ) , axis = 0) - if node._depth_out is None: - node._depth_out = np.sum([node.result._depth_out*w for n,w in zip(interpolationNodes[0:1], [1] )], axis=0) - if node.temperature_out is None: - node.temperature_out = np.sum([n.result.temperature_out*w for n,w in zip(interpolationNodes[0:1], [1] )], axis=0) + # if node._depth_out is None: + # node._depth_out = np.sum([node.result._depth_out*w for n,w in zip(interpolationNodes[0:1], [1] )], axis=0) + # if node.temperature_out is None: + # node.temperature_out = np.sum([n.result.temperature_out*w for n,w in zip(interpolationNodes[0:1], [1] )], axis=0) if node.sed is None: node.sed = np.sum([n.sed*w for n,w in zip(interpolationNodes,iWeightNorm)], axis=0) @@ -167,7 +175,7 @@ def interpolateNode(interpolationNodes: List[single_node], interpolationWeights= def interpolate_all_nodes(builder:Builder)->Builder: for ni in range(len(builder.nodes)): for nj in range(len(builder.nodes[ni])): - if builder.nodes[ni][nj] is False: + if (builder.nodes[ni][nj] is False) or (not builder.nodes[ni][nj]._full_simulation): closest_x_up = [] for j in range(ni,len(builder.nodes[nj])): matching_x = [ i[0] for i in builder.indexer_full_sim if i[0]==j ] @@ -195,6 +203,7 @@ def interpolate_all_nodes(builder:Builder)->Builder: interpolationNodes = [ builder.nodes[i[0]][i[1]] for i in itertools.product(closest_x_up+closest_x_down, closest_y_up+closest_y_down) ] interpolationNodes = [nn for nn in interpolationNodes if nn is not False] + interpolationNodes = [nn for nn in interpolationNodes if nn._full_simulation] node = interpolateNode(interpolationNodes) node.X, node.Y = builder.grid.location_grid[ni,nj,:] builder.nodes[ni][nj] = node diff --git a/warmth/postprocessing.py b/warmth/postprocessing.py index 435cb18..de8af1c 100644 --- a/warmth/postprocessing.py +++ b/warmth/postprocessing.py @@ -9,7 +9,8 @@ class Results: """Simulation results """ - def __init__(self,depth:np.ndarray, temperature:np.ndarray,sediments_ids:np.ndarray,sediment_input:pd.DataFrame,k_crust:float,k_lith:float,k_asth:float): + def __init__(self, depth:np.ndarray, temperature:np.ndarray, + sediments_ids:np.ndarray,sediment_input:pd.DataFrame,k_crust:float,k_lith:float,k_asth:float): self._depth=depth self._temperature=temperature self._sediments_ids=sediments_ids @@ -22,7 +23,7 @@ class resultValues(TypedDict): depth: np.ndarray[np.float64] layerId: np.ndarray[np.int32] values:np.ndarray[np.float64] - + @property def ages(self)->np.ndarray[np.int32]: """Array of all simulated ages @@ -32,7 +33,7 @@ def ages(self)->np.ndarray[np.int32]: np.ndarray Array of ages """ - return np.arange(self._depth.shape[1],dtype=np.int32) + return np.arange(self._max_time,dtype=np.int32) def top_crust(self,age:int)->float: """Depth of crust @@ -62,7 +63,7 @@ def top_lithosphere(self,age:int)->float: ------- float Depth of lithospheric mantle / Moho from sea level (m) - """ + """ depth_idx= np.where(self.sediment_ids(age) == -2)[0][0] return self._depth[depth_idx,age] @@ -78,7 +79,7 @@ def top_asthenosphere(self,age:int)->float: ------- float Depth of Asthenosphere from sea level (m) - """ + """ depth_idx= np.where(self.sediment_ids(age) == -3)[0][0] return self._depth[depth_idx,age] @@ -572,7 +573,7 @@ def interp_value(self): for n in self._builder.iter_node(): if n._full_simulation is False: idx = n.indexer - val =interped[idx[0],idx[1]] + val =interped[idx[1],idx[0]] setattr(n,prop,val) return @@ -594,7 +595,7 @@ def interp_arr(self): if isinstance(getattr(node,prop),type(None)): setattr(node,prop,np.zeros(self.n_age)) idx = node.indexer - interpolated_val =interp_all_this_age[idx[0],idx[1]] + interpolated_val =interp_all_this_age[idx[1],idx[0]] arr = getattr(node,prop) arr[age] =interpolated_val setattr(node,prop,arr) diff --git a/warmth/simulator.py b/warmth/simulator.py index d069c46..6be4565 100644 --- a/warmth/simulator.py +++ b/warmth/simulator.py @@ -53,6 +53,9 @@ def run(self) -> Path: self.node = fw.current_node self._pad_sediments() self.node.simulated_at = time.time() + self.node.max_time = self.node._depth_out.shape[1] + self.node.compute_derived_arrays() + self.node.node_path = self.node_path filepath = self._save_results() # Delete input node self.node_path.unlink(missing_ok=True) @@ -244,12 +247,13 @@ def _parellel_run(self, save, purge, use_mpi=False): if save==False: from shutil import rmtree rmtree(self._builder.parameters.output_path) - if filtered >0: - logger.info(f"Interpolating results back to {filtered} partial simulated nodes") - interp_res= Results_interpolator(self._builder) - interp_res.run() + # if filtered >0: + # logger.info(f"Interpolating results back to {filtered} partial simulated nodes") + # interp_res= Results_interpolator(self._builder) + # interp_res.run() return def put_node_to_grid(self,node:single_node): + node.clear_unused_data() self._builder.nodes[node.indexer[0]][node.indexer[1]]=node return