-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathget_point.py
183 lines (153 loc) · 8.63 KB
/
get_point.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
import h5py
import numpy as np
import typer
from scipy.spatial import KDTree
from typing_extensions import Annotated
from generating_ASAGI_file import writeNetcdf4SeisSol, writeNetcdf4Paraview
from netCDF4 import Dataset
from meshgen import get_cartesian, apply_centering_points, apply_rotation_points
from tqdm import tqdm
def createNetcdf4ParaviewHandle(sname, x, y, z, aName):
"create a netcdf file readable by paraview (but not by ASAGI)"
fname = sname + "_paraview.nc"
# print("writing " + fname)
# Creating the netcdf file
nx = x.shape[0]
ny = y.shape[0]
nz = z.shape[0]
rootgrp = Dataset(fname, "w", format="NETCDF4")
rootgrp.createDimension("u", nx)
rootgrp.createDimension("v", ny)
rootgrp.createDimension("w", nz)
vx = rootgrp.createVariable("u", "f4", ("u",))
vx[:] = x
vy = rootgrp.createVariable("v", "f4", ("v",))
vy[:] = y
vz = rootgrp.createVariable("w", "f4", ("w",))
vz[:] = z
vTd = rootgrp.createVariable(aName, "f4", ("u", "v", "w"))
return rootgrp, vTd
def createNetcdf4SeisSolHandle(sname, x, y, z, aName):
"create a netcdf file readable by paraview (but not by ASAGI)"
fname = sname + "_ASAGI.nc"
# print("writing " + fname)
# Creating the netcdf file
nx = x.shape[0]
ny = y.shape[0]
nz = z.shape[0]
rootgrp = Dataset(fname, "w", format="NETCDF4")
rootgrp.createDimension("u", nx)
rootgrp.createDimension("v", ny)
rootgrp.createDimension("w", nz)
vx = rootgrp.createVariable("u", "f4", ("u",))
vx[:] = x
vy = rootgrp.createVariable("v", "f4", ("v",))
vy[:] = y
vz = rootgrp.createVariable("w", "f4", ("w",))
vz[:] = z
vTd = rootgrp.createVariable("data", "f4", ("u", "v", "w"))
return rootgrp, vTd
def main(
meta_file: Annotated[str, typer.Argument(help="Path for the meta file generated by meshgen script")],
lat: Annotated[str, typer.Option(help="Latitude to convert to point on the mesh")] = 37.341206616740266,
long: Annotated[str, typer.Option(help="Longitude to convert to point on the mesh")] = -121.88075569799896,
depth: Annotated[
float, typer.Option(help="How deep should the point be from the sea level (in meters)")] = -1000.0,
# from sea level in meters
round_to_closest_point: Annotated[
bool, typer.Option(help="Should the point be rounded to the nearest point on the fault line")] = True,
point_field_resolution: Annotated[
float, typer.Option(help="Resolution for netcdf files (in m)")] = 100,
output_distance_from_faults: Annotated[
str, typer.Option(help="Generate netcdf file for distance from fault points")] = None,
output_distance_from_topo: Annotated[
str, typer.Option(help="Generate netcdf file for distance from topography points")] = None,
chunk_size: Annotated[
int, typer.Option(help="Size of chunks while generating")] = 50
):
with h5py.File(meta_file, "r") as f:
closest_point = np.array([long, lat, 0])
if round_to_closest_point:
all_long_lats = f["all_long_lats"][:]
center = f["center"][:]
rotation_matrix = f["rotation_matrix"][:]
tree = KDTree(all_long_lats[:, np.array([1, 0])])
_, ii = tree.query([[lat, long], ], k=[1])
closest_point = all_long_lats[ii[0][0],]
print(f"Rounding to lat,long on the mesh {closest_point[1]} {closest_point[0]}")
cart = get_cartesian(lat_deg=closest_point[1], lon_deg=closest_point[0], alt=depth)
cart = apply_rotation_points(cart, rotation_matrix)
cart = apply_centering_points(cart, center)
print(f"Location of the point is {cart}")
if output_distance_from_faults is not None or output_distance_from_topo is not None:
if f.get("bounding_box") is None:
print("Could not find bounding box in meta file, create bounding box with mesh")
exit()
print("generating bounding box")
bounding_box = f.get("bounding_box")[:]
min_coords = np.min(bounding_box, axis=0)
max_coords = np.max(bounding_box, axis=0)
x = np.linspace(min_coords[0], max_coords[0], int((max_coords[0] - min_coords[0]) / point_field_resolution))
y = np.linspace(min_coords[1], max_coords[1], int((max_coords[0] - min_coords[0]) / point_field_resolution))
z = np.linspace(min_coords[2], max_coords[2], int((max_coords[0] - min_coords[0]) / point_field_resolution))
xg, yg, zg = np.meshgrid(x, y, z, indexing='ij')
if output_distance_from_faults is not None:
print("generating fault KDTree")
fault_tree = KDTree(f["fault_points"][:])
rootgrp, vTd = createNetcdf4ParaviewHandle(output_distance_from_faults, z, y, x, "fault_distance")
rootgrp_ss, vTd_ss = createNetcdf4SeisSolHandle(output_distance_from_faults, z, y, x, "fault_distance")
total_iterations = (len(x) // chunk_size + 1) * (len(y) // chunk_size + 1) * (len(z) // chunk_size + 1)
progress_bar = tqdm(total=total_iterations, desc="Generating fault_distance")
# Loop over the grid in chunks
for i in range(0, len(x), chunk_size):
for j in range(0, len(y), chunk_size):
for k in range(0, len(z), chunk_size):
x_chunk = x[i:i + chunk_size]
y_chunk = y[j:j + chunk_size]
z_chunk = z[k:k + chunk_size]
xg, yg, zg = np.meshgrid(x_chunk, y_chunk, z_chunk, indexing='ij')
# Query the fault tree
sub_distances, index = fault_tree.query(
np.stack((xg.flatten(), yg.flatten(), zg.flatten())).T, k=[1])
sub_distances = sub_distances.squeeze().reshape(xg.shape)
sub_distances = np.einsum('ijk->kji', sub_distances)
# Store the results in the appropriate slice of the distances array
vTd[k:k + chunk_size, j:j + chunk_size, i:i + chunk_size] = sub_distances
vTd_ss[k:k + chunk_size, j:j + chunk_size, i:i + chunk_size] = sub_distances
# Update tqdm progress bar
progress_bar.update(1)
# Close tqdm progress bar
progress_bar.close()
rootgrp.close()
rootgrp_ss.close()
if output_distance_from_topo is not None:
print("generating topography KDTree")
topo_tree = KDTree(f["topo_points"][:])
rootgrp, vTd = createNetcdf4ParaviewHandle(output_distance_from_topo, z, y, x, "topo_distance")
rootgrp_ss, vTd_ss = createNetcdf4SeisSolHandle(output_distance_from_topo, z, y, x, "topo_distance")
total_iterations = (len(x) // chunk_size + 1) * (len(y) // chunk_size + 1) * (len(z) // chunk_size + 1)
progress_bar = tqdm(total=total_iterations, desc="Generating topo_distance")
for i in range(0, len(x), chunk_size):
for j in range(0, len(y), chunk_size):
for k in range(0, len(z), chunk_size):
x_chunk = x[i:i + chunk_size]
y_chunk = y[j:j + chunk_size]
z_chunk = z[k:k + chunk_size]
xg, yg, zg = np.meshgrid(x_chunk, y_chunk, z_chunk, indexing='ij')
# Query the fault tree
sub_distances, index = topo_tree.query(
np.stack((xg.flatten(), yg.flatten(), zg.flatten())).T, k=[1])
sub_distances = sub_distances.squeeze().reshape(xg.shape)
sub_distances = np.einsum('ijk->kji', sub_distances)
# Store the results in the appropriate slice of the distances array
vTd[k:k + chunk_size, j:j + chunk_size, i:i + chunk_size] = sub_distances
vTd_ss[k:k + chunk_size, j:j + chunk_size, i:i + chunk_size] = sub_distances
# Update tqdm progress bar
progress_bar.update(1)
# Close tqdm progress bar
progress_bar.close()
rootgrp.close()
rootgrp_ss.close()
if __name__ == "__main__":
# main("outputs/meta.h5", 37.341206616740266, -121.88075569799896, -1000)
typer.run(main)