Skip to content

Consistent coordinates for pre-compute and 1d profile generation #50

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sungeunbae merged 2 commits into from
Feb 25, 2026
Merged

Consistent coordinates for pre-compute and 1d profile generation #50

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
572e265
Consistent coordinates for pre-compute and 1d profile generation
sungeunbae Feb 23, 2026
9a981e6
Update velocity_modelling/scripts/generate_1d_profiles.py
sungeunbae Feb 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 29 additions & 14 deletions velocity_modelling/geometry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -224,7 +224,11 @@ def copy(self) -> Self:

@njit
def point_on_vertex(
boundary_lats: np.ndarray, boundary_lons: np.ndarray, lat: float, lon: float
boundary_lats: np.ndarray,
boundary_lons: np.ndarray,
lat: float,
lon: float,
tol: float = 1e-07,
) -> bool:
"""
Check if a point lies on a vertex of a basin boundary.
Expand All @@ -239,17 +243,20 @@ def point_on_vertex(
Latitude of the point to check.
lon : float
Longitude of the point to check.
tol : float
Tolerance for floating point equality test, default: 1e-07


Returns
-------
bool
True if the point matches a boundary vertex within tolerance (1e-07), otherwise False.
True if the point matches a boundary vertex within tolerance, otherwise False.
"""
# Vectorized comparison using NumPy with tolerance
lat_matches = (
np.abs(boundary_lats - lat) <= 1e-07
np.abs(boundary_lats - lat) <= tol
) # np.isclose is not supported by njit
lon_matches = np.abs(boundary_lons - lon) <= 1e-07
lon_matches = np.abs(boundary_lons - lon) <= tol

matches = np.logical_and(lat_matches, lon_matches)
Comment on lines 225 to 261

@lispandfound lispandfound Feb 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is just np.isclose. both np.isclose and np.allclose have been supported by numba since four years ago, also here.

lat_close = np.isclose(lat, boundary_lats, rtol=0.0, atol=tol)
lon_close = np.isclose(lon, boundary_lons, rtol=0.0, atol=tol)
return np.any(lat_close & lon_close)
>>> @numba.njit
... def test(lat, lon, boundary_lats, boundary_lons, tol=1e-7):
...     lat_close = np.isclose(lat, boundary_lats, rtol=0.0, atol=tol)
...     lon_close = np.isclose(lon, boundary_lons, rtol=0.0, atol=tol)
...     return np.any(lat_close & lon_close)
...
>>> test(-43.0, 172.64, np.array([-43.0, -44.0]), np.array([172.0, 172.65]))
False
>>> test(-43.0, 172.64, np.array([-43.0, -44.0]), np.array([172.0, 172.65]))
False
>>> test(-43.0, 172.64, np.array([-43.0, -44.0]), np.array([172.0, 172.65]))
False
>>> test(-43.0, 172.64, np.array([-43.0, -44.0]), np.array([172.0, 172.64]))
False
>>> test(-43.0, 172.64, np.array([-43.0, -43.0]), np.array([172.0, 172.64]))
True

@sungeunbae sungeunbae Feb 24, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@lispandfound Thanks for the update regarding np.isclose() support in Numba—I wasn't aware it was supported now.However, after reviewing the documentation, I prefer to stick with np.abs(a - b) <= tol for this specific case.

np.isclose uses a relative tolerance by default (rtol=1e-05)

numpy.isclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False)

and internally checks

 absolute(a - b) <= (atol + rtol * absolute(b))

For geographic coordinates, we need a fixed spatial tolerance (e.g., 1e-7 degrees) regardless of the coordinate magnitude (e.g., relative error at 170° is much larger than at 1°). To achieve this with np.isclose(), I would have to explicitly disable the relative tolerance

np.isclose(a, b, rtol=0, atol=tol)

This is more verbose, and less readable than the direct comparision.. Also inside @njit, direct subtraction is the most efficient primitive operation avoiding the extra overheads of np.isclose() (eg. internal handing of NaN, infinites etc).

https://numpy.org/doc/stable/reference/generated/numpy.isclose.html

@sungeunbae sungeunbae Feb 24, 2026 •
edited
Loading

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Based on my observation, I even feel like to replace the 2 uses of "np.isclose()" with np.abs() <=tol

Line 1032-1040

    lat = np.where(
        np.isclose(zg, 0),
        0,
        90 - np.arctan(np.sqrt(xg**2 + yg**2) / zg) / RPERD - np.where(zg < 0, 180, 0),
    )

    lon = np.where(
        np.isclose(xg, 0), 0, np.arctan(yg / xg) / RPERD - np.where(xg < 0, 180, 0)
    )

Perhaps this is not critical as it is comparing with 0, and np.isclose(x,0) is abs(x-0) <= atol + rtol*0, effectively `abs(x)<=atol (which is 1e-8)

But for the consistency, I feel like to go ahead. Thought? @lispandfound


Expand Down Expand Up @@ -335,7 +342,9 @@ def find_edge_inds(


@njit
def find_corner_inds(lats: np.ndarray, lons: np.ndarray, lat: float, lon: float):
def find_corner_inds(
lats: np.ndarray, lons: np.ndarray, lat: float, lon: float, tol: float = 1e-07
):
"""
Find the corner indices of the given latitude and longitude arrays.

Expand All @@ -349,6 +358,8 @@ def find_corner_inds(lats: np.ndarray, lons: np.ndarray, lat: float, lon: float)
Latitude of the given point.
lon : float
Longitude of the given point.
tol : float
Tolerance for floating point equality test, default : 1e-07

Returns
-------
Expand All @@ -361,18 +372,18 @@ def find_corner_inds(lats: np.ndarray, lons: np.ndarray, lat: float, lon: float)
nlat = len(lats)
nlon = len(lons)

if np.abs(lat - lats[0]) <= 1e-07:
if np.abs(lat - lats[0]) <= tol:
corner_lat_ind = 0
elif np.abs(lat - lats[-1]) <= 1e-07:
elif np.abs(lat - lats[-1]) <= tol:
corner_lat_ind = nlat - 1
else:
raise ValueError(
f"Point lies outside of surface bounds. Lat {lat} outside [{lats[0]} {lats[-1]}]"
)

if np.abs(lon - lons[0]) <= 1e-07:
if np.abs(lon - lons[0]) <= tol:
corner_lon_ind = 0
elif np.abs(lon - lons[-1]) <= 1e-07:
elif np.abs(lon - lons[-1]) <= tol:
corner_lon_ind = nlon - 1
Comment on lines -364 to 387

@lispandfound lispandfound Feb 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similarly here you should use np.isclose

else:
raise ValueError(
Expand All @@ -384,7 +395,7 @@ def find_corner_inds(lats: np.ndarray, lons: np.ndarray, lat: float, lon: float)

@njit
def find_basin_adjacent_points_numba(
lats: np.ndarray, lons: np.ndarray, lat: float, lon: float
lats: np.ndarray, lons: np.ndarray, lat: float, lon: float, tol: float = 1e-07
) -> tuple[bool, np.ndarray, np.ndarray]:
"""
Identify adjacent latitude/longitude indices for a given point in basin coordinates.
Expand All @@ -399,6 +410,9 @@ def find_basin_adjacent_points_numba(
Target latitude for adjacency checks.
lon : float
Target longitude for adjacency checks.
tol : float
Tolerance for floating point equality test, default: 1e-7


Returns
-------
Expand All @@ -424,10 +438,11 @@ def find_basin_adjacent_points_numba(
if 0 < lat_idx < len(lats):
lat_ind[0] = lat_idx - 1
lat_ind[1] = lat_idx
elif lat_idx == 0 and lats[0] == lat:

elif lat_idx == 0 and np.abs(lats[0] - lat) <= tol:
lat_ind[0] = 0
lat_ind[1] = 1
elif lat_idx == len(lats) and lats[-1] == lat:
elif lat_idx == len(lats) and np.abs(lats[-1] - lat) <= tol:
lat_ind[0] = len(lats) - 1
lat_ind[1] = len(lats)

Expand All @@ -441,10 +456,10 @@ def find_basin_adjacent_points_numba(
if 0 < lon_idx < len(lons):
lon_ind[0] = lon_idx - 1
lon_ind[1] = lon_idx
elif lon_idx == 0 and lons[0] == lon:
elif lon_idx == 0 and np.abs(lons[0] - lon) <= tol:
lon_ind[0] = 0
lon_ind[1] = 1
elif lon_idx == len(lons) and lons[-1] == lon:
elif lon_idx == len(lons) and np.abs(lons[-1] - lon) <= tol:
lon_ind[0] = len(lons) - 1
lon_ind[1] = len(lons)

Expand Down
75 changes: 48 additions & 27 deletions velocity_modelling/scripts/generate_1d_profiles.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -542,29 +542,34 @@ def generate_1d_profiles(
# PRE-COMPUTE BASIN MEMBERSHIP FOR ALL PROFILES
logger.log(logging.INFO, "Pre-computing basin membership for all profiles")

basin_membership = BasinMembership(
basin_data_list,
smooth_boundary=nz_tomography_data.smooth_boundary,
logger=logger,
)
profile_basin_membership = basin_membership.check_stations(
df["lat"].values, df["lon"].values
)

logger.log(logging.INFO, f"Basin membership computed for {len(df)} profiles")

for i in tqdm(range(len(df)), desc="Generating profiles", unit="profile"):
logger.log(logging.INFO, f"Generating profile {i + 1} of {len(df)}")

# Pre-calculate mesh coordinates to ensure consistency with basin checks.
# This prevents "point outside basin" errors caused by slight shifts between
# input CSV coordinates and the generated model grid, while maintaining
# the speed of vectorized basin checks.
mesh_lats = np.zeros(len(df))
mesh_lons = np.zeros(len(df))
global_meshes = []

# Use a silent logger to avoid spamming logs during coordinate generation
silent_logger = logging.getLogger("silent_mesh_gen")
silent_logger.setLevel(logging.WARNING)

lats = df["lat"].values
lons = df["lon"].values
zmins = df["zmin"].values
zmaxs = df["zmax"].values
spacings = df["spacing"].values

for i in tqdm(range(len(df)), desc="Pre-computing meshes"):
# Initialize model extent
model_extent = {
"version": vm_params["model_version"],
"origin_rot": 0.0,
"extent_x": 1.0,
"extent_y": 1.0,
"h_lat_lon": 1.0,
"origin_lat": df["lat"].iloc[i],
"origin_lon": df["lon"].iloc[i],
"origin_lat": lats[i],
"origin_lon": lons[i],
}

# Set depth parameters
Expand All @@ -574,13 +579,10 @@ def generate_1d_profiles(
model_extent["h_depth"] = 1.0 # Placeholder, as actual depths are set later
else:
spacing_offset = 0.5
model_extent["extent_zmin"] = (
df["zmin"].iloc[i] - spacing_offset * df["spacing"].iloc[i]
)
model_extent["extent_zmax"] = (
df["zmax"].iloc[i] + spacing_offset * df["spacing"].iloc[i]
)
model_extent["h_depth"] = df["spacing"].iloc[i]
model_extent["extent_zmin"] = zmins[i] - spacing_offset * spacings[i]
model_extent["extent_zmax"] = zmaxs[i] + spacing_offset * spacings[i]
model_extent["h_depth"] = spacings[i]

model_extent["nx"] = int(
model_extent["extent_x"] / model_extent["h_lat_lon"] + 0.5
)
Expand All @@ -593,15 +595,34 @@ def generate_1d_profiles(
+ 0.5
)

# Generate mesh (only for depth calculation, not basin membership)
global_mesh = gen_full_model_grid_great_circle(model_extent, logger)
gm = gen_full_model_grid_great_circle(model_extent, silent_logger)

if depth_values:
gm.nz = len(depth_values)
gm.z = np.array([-1000 * dep for dep in depth_values])

global_meshes.append(gm)
mesh_lats[i] = gm.lat.flat[0]
mesh_lons[i] = gm.lon.flat[0]

basin_membership = BasinMembership(
basin_data_list,
smooth_boundary=nz_tomography_data.smooth_boundary,
logger=logger,
)
profile_basin_membership = basin_membership.check_stations(mesh_lats, mesh_lons)

logger.log(logging.INFO, f"Basin membership computed for {len(df)} profiles")

for i in tqdm(range(len(df)), desc="Generating profiles", unit="profile"):
logger.log(logging.INFO, f"Generating profile {i + 1} of {len(df)}")

global_mesh = global_meshes[i]
if depth_values:
global_mesh.nz = len(depth_values)
logger.log(
logging.DEBUG,
f"Number of model points - nx: {global_mesh.nx}, ny: {global_mesh.ny}, nz: {global_mesh.nz}",
)
global_mesh.z = np.array([-1000 * dep for dep in depth_values])

basin_indices = profile_basin_membership[i]

Expand Down