|
| 1 | + |
| 2 | +from mpi4py import MPI |
| 3 | +import json |
| 4 | +import os |
| 5 | +import numpy as np |
| 6 | +import landlab |
| 7 | +from landlab.components import AdvectionSolverTVD |
| 8 | + |
| 9 | +from landlab.io.legacy_vtk import write_legacy_vtk |
| 10 | + |
| 11 | +current_time = 0 |
| 12 | + |
| 13 | +comm = None |
| 14 | + |
| 15 | +model_grid = None |
| 16 | +elevation = None |
| 17 | +surface_advector = None |
| 18 | +horizontal_velocity = None |
| 19 | +ux = None |
| 20 | +uy = None |
| 21 | + |
| 22 | +timestep = 0 |
| 23 | +vtks = [] |
| 24 | + |
| 25 | +def initialize(comm_handle): |
| 26 | + if not comm_handle is None: |
| 27 | + # Convert the handle back to an MPI communicator |
| 28 | + global comm |
| 29 | + comm = MPI.Comm.f2py(comm_handle) |
| 30 | + |
| 31 | + rank = comm.Get_rank() |
| 32 | + size = comm.Get_size() |
| 33 | + |
| 34 | + print(f"Python: Hello from Rank {rank} of {size}") |
| 35 | + |
| 36 | + data = 1 |
| 37 | + globalsum = comm.allreduce(data, op=MPI.SUM) |
| 38 | + if comm.rank == 0: |
| 39 | + print(f"\tPython: testing communication; sum {globalsum}") |
| 40 | + else: |
| 41 | + print("Python: running sequentially!") |
| 42 | + |
| 43 | +def finalize(): |
| 44 | + pass |
| 45 | + |
| 46 | + |
| 47 | +# Run the Landlab simulation from the current time to end_time and return |
| 48 | +# the new topographic elevation (in m) at each local node. |
| 49 | +# dict_variable_name_to_value_in_nodes is a dictionary mapping variables |
| 50 | +# (x velocity, y velocity, temperature, etc.) to an array of values in each |
| 51 | +# node. |
| 52 | +def update_until(end_time, dict_variable_name_to_value_in_nodes): |
| 53 | + global current_time, elevation, surface_advector, horizontal_velocity, ux, uy, timestep |
| 54 | + dt = end_time - current_time |
| 55 | + timestep += 1 |
| 56 | + |
| 57 | + deposition_erosion = np.zeros(model_grid.number_of_nodes) |
| 58 | + |
| 59 | + x_velocity = dict_variable_name_to_value_in_nodes["x velocity"] |
| 60 | + y_velocity = dict_variable_name_to_value_in_nodes["y velocity"] |
| 61 | + z_velocity = dict_variable_name_to_value_in_nodes["z velocity"] |
| 62 | + |
| 63 | + # The x, y, z velocities passed in from ASPECT are defined at the node points of the LandLab mesh. |
| 64 | + # However, lateral advection of the LandLab surface requires that the velocity is defined at the links. |
| 65 | + # To do this, we determine the average value of the x,y velocity at the link using the node points that |
| 66 | + # are connected by the link. Then, we can use AdvectionSolverTVD to advect the surface. |
| 67 | + |
| 68 | + x_vel_at_links = model_grid.map_mean_of_link_nodes_to_link(x_velocity) |
| 69 | + y_vel_at_links = model_grid.map_mean_of_link_nodes_to_link(y_velocity) |
| 70 | + |
| 71 | + horizontal_velocity[model_grid.horizontal_links] = x_vel_at_links[model_grid.horizontal_links] |
| 72 | + horizontal_velocity[model_grid.vertical_links] = y_vel_at_links[model_grid.vertical_links] |
| 73 | + surface_advector = AdvectionSolverTVD(model_grid, fields_to_advect=elevation) |
| 74 | + |
| 75 | + # Substepping for surface processes |
| 76 | + if dt>0: |
| 77 | + n_substeps = 10 |
| 78 | + sub_dt = dt / n_substeps |
| 79 | + for _ in range(n_substeps): |
| 80 | + |
| 81 | + # TODO: |
| 82 | + elevation_before = elevation.copy() |
| 83 | + |
| 84 | + # Uplift the topography using the ASPECT vertical velocity. |
| 85 | + elevation[model_grid.core_nodes] += z_velocity[model_grid.core_nodes] * sub_dt |
| 86 | + |
| 87 | + # Advect the grid with the horizontal velocity from ASPECT. The ASPECT velocities are already being passed |
| 88 | + # correctly at each NODE point of the landlab grid. However, the TVDAdvector requires the velocity at the LINKS. |
| 89 | + # There must be a relationship between the two that can be leveraged, but at the moment I'm not sure how. |
| 90 | + surface_advector.update(sub_dt) |
| 91 | + |
| 92 | + deposition_erosion += elevation - elevation_before |
| 93 | + |
| 94 | + current_time = end_time |
| 95 | + |
| 96 | + # filename = f"./output/landlab_{str(timestep).zfill(3)}.vtk" |
| 97 | + # print("Writing output VTK file...", filename) |
| 98 | + # vtk_file = write_legacy_vtk(path=filename, grid=model_grid, clobber=True) |
| 99 | + # vtks.append((current_time, filename)) |
| 100 | + |
| 101 | + # if True: |
| 102 | + # # write vtk.series file (ParaView supports legacy VTK in this format) |
| 103 | + # with open("./output/landlab.vtk.series", "w") as f: |
| 104 | + # series = { |
| 105 | + # "file-series-version": "1.0", |
| 106 | + # "files": [ |
| 107 | + # {"name": os.path.basename(filename), "time": time} |
| 108 | + # for time, filename in vtks |
| 109 | + # ] |
| 110 | + # } |
| 111 | + # json.dump(series, f, indent=2) |
| 112 | + |
| 113 | + # print ("deposition/erosion:", np.linalg.norm(deposition_erosion)) |
| 114 | + print("Max elevation:", np.max(elevation), "Min elevation:", np.min(elevation)) |
| 115 | + # print("Max elevation:", np.max(deposition_erosion), "Min elevation:", np.min(deposition_erosion)) |
| 116 | + |
| 117 | + return deposition_erosion |
| 118 | + # return elevation |
| 119 | + |
| 120 | +def set_mesh_information(dict_grid_information): |
| 121 | + global model_grid, elevation, horizontal_velocity, ux, uy |
| 122 | + |
| 123 | + if not model_grid: |
| 124 | + print("* Creating RasterModelGrid ...") |
| 125 | + x_extent = dict_grid_information["Mesh X extent"] |
| 126 | + y_extent = dict_grid_information["Mesh Y extent"] |
| 127 | + spacing = dict_grid_information["Mesh Spacing"] |
| 128 | + |
| 129 | + nrows = int(y_extent / spacing) # number of node rows |
| 130 | + ncols = int(x_extent / spacing) # number of node columns |
| 131 | + |
| 132 | + model_grid = landlab.RasterModelGrid((nrows, ncols), xy_spacing=(spacing, spacing)) |
| 133 | + model_grid.set_closed_boundaries_at_grid_edges(True, True, True, True) |
| 134 | + |
| 135 | + print("* Creating topographic elevation ...") |
| 136 | + # Initialize topography array with zeros |
| 137 | + elevation = model_grid.add_zeros("topographic__elevation", at="node") |
| 138 | + |
| 139 | + # create a Gaussian hill in the center of the domain |
| 140 | + gaussian_height = 5e3 |
| 141 | + x_extent = model_grid.x_of_node.max() - model_grid.x_of_node.min() |
| 142 | + y_extent = model_grid.y_of_node.max() - model_grid.y_of_node.min() |
| 143 | + elevation += gaussian_height * np.exp(-((model_grid.node_x - x_extent/2)**2 + (model_grid.node_y - y_extent/2)**2) / (2 * (10e3)**2)) |
| 144 | + |
| 145 | + # Initialize the horizontal velocity so that we can pass the ASPECT velocities to it later. |
| 146 | + horizontal_velocity = model_grid.add_zeros("advection__velocity", at="link") |
| 147 | + |
| 148 | + ux = model_grid.add_zeros("x_velocity", at="node") |
| 149 | + uy = model_grid.add_zeros("y_velocity", at="node") |
| 150 | + print("\tnumber of nodes:", model_grid.number_of_nodes) |
| 151 | + |
| 152 | + print("* Done") |
| 153 | + |
| 154 | +# Return the x coordinates of the locally owned nodes on this |
| 155 | +# MPI rank. grid_id is always 0. |
| 156 | +def get_grid_x(grid_id): |
| 157 | + global model_grid |
| 158 | + return model_grid.node_x |
| 159 | + |
| 160 | +# Return the y coordinates of the locally owned nodes on this |
| 161 | +# MPI rank. grid_id is always 0. |
| 162 | +def get_grid_y(grid_id): |
| 163 | + global model_grid |
| 164 | + return model_grid.node_y |
| 165 | + |
| 166 | +def initalize_landlab_components(landlab_component_parameters): |
| 167 | + global model_grid, elevation |
| 168 | + |
| 169 | + pass |
| 170 | + |
| 171 | +# Return the initial topography at the start of the simulation |
| 172 | +# in each node. |
| 173 | +def get_initial_topography(grid_id): |
| 174 | + global elevation |
| 175 | + return elevation |
| 176 | + |
| 177 | + |
| 178 | +def write_output(timestep): |
| 179 | + # Write the grid to vtk |
| 180 | + print("Writing output VTK file...") |
| 181 | + vtk_file = write_legacy_vtk(path=f"./output_{str(timestep).zfill(3)}.vtk", grid=model_grid, clobber=True) |
0 commit comments