Collocation
This section provides the core tools for collocating model and observational data.
Core Collocation
The main Collocate class is the primary entry point for running the collocation analysis.
!!! note
The Collocate class is designed to be used with a Model object and an Observation object as inputs.
ocstrack.Collocation.collocate
Collocate Module
Collocate
Model–observation collocation engine (2D and 3D).
This is the main class. It handles the spatial and temporal collocation of observation data (Satellite or Argo) with unstructured model outputs.
It dispatches to either a 2D/surface ('2D', '3D_Surface') or 3D/profile ('3D_Profile') collocation strategy based on the model_run's 'var_type'.
Methods
run(output_path=None) Run the full collocation and return a combined xarray.Dataset.
__init__(model_run, observation, dist_coast=None, n_nearest=4, search_radius=None, time_buffer=None, weight_power=1.0, temporal_interp=False)
Initialize the Collocate object.
Parameters
model_run : object Initialized Model object containing grid and file info. observation : Union[SatelliteData, ArgoData] Initialized Observation object (e.g., SatelliteData or ArgoData). dist_coast : xarray.Dataset, optional Optional dataset containing distance-to-coast info. n_nearest : int, optional Number of nearest spatial model nodes to use (default=4). search_radius : float, optional Radius (in meters) to search for spatial neighbors. If provided, overwrites n_nearest. time_buffer : np.timedelta64, optional Temporal search buffer; if None, inferred from model timestep. weight_power : float, default=1.0 Power exponent for inverse distance weighting. temporal_interp : bool, default=False Whether to perform linear temporal interpolation (default=False).
Raises
ValueError If both or neither of 'n_nearest'/'search_radius' are provided, or if time buffer cannot be inferred.
run(output_path=None)
Run the full model–observation collocation process.
Dispatches to the correct method (surface or profile) based on the 'var_type' provided during initialization.
Parameters
output_path : str, optional If provided, writes collocated output to this NetCDF file path.
Returns
xarray.Dataset Dataset containing collocated observation and model data.
Raises
TypeError If the observation object type does not match the collocation type (e.g., using ArgoData for '2D' collocation). NotImplementedError If the 'var_type' is not recognized.
Output & Spatial/Temporal Helpers
These modules provide helper functions for handling the output of the collocation, as well as for spatial and temporal operations.
Output
ocstrack.Collocation.output
Functions for writing collocated outputs file
get_max_neighbors(result_list)
Get the maximum number of neighbors from a list of 2D arrays.
Used to handle ragged results from radius search.
Parameters
result_list : list A list of 2D numpy arrays (n_obs, n_neighbors)
Returns
int The maximum number of neighbors (max columns) found in the list.
make_collocated_nc_2d(results, n_nearest=None, model_var_name='model_var', obs_var_name='obs_var')
Format the 2D collocated data into a CF-compliant xarray Dataset.
Parameters
results : dict Dictionary containing the concatenated lists of collocated data. n_nearest : int, optional The number of nearest neighbors used (k). If None, inferred from the radius search results. model_var_name : str, default='model_var' The name of the model variable (e.g., 'sigWaveHeight'). obs_var_name : str, default='obs_var' The name of the observation variable (e.g., 'swh').
Returns
xarray.Dataset A CF-compliant dataset of the collocated 2D surface data.
make_collocated_nc_3d(results, max_levels)
Format the 3D collocated data into a CF-compliant xarray Dataset.
Parameters
results : dict Dictionary containing the final 3D profile arrays. max_levels : int The size of the padded vertical dimension ('n_levels').
Returns
xarray.Dataset A CF-compliant dataset of the collocated 3D profile data.
pad_arrays_to_max(arrays, max_cols)
Pad 2D arrays with NaNs to have the same number of columns.
Used to stack ragged results from radius search into a single 2D array.
Parameters
arrays : list List of 2D numpy arrays (n_obs_i, n_neighbors_i). max_cols : int The target number of columns (max neighbors) to pad to.
Returns
np.ndarray A single, stacked 2D array of shape (N_total_obs, max_cols).
Spatial
ocstrack.Collocation.spatial
Functions for spatial collocation
GeocentricSpatialLocator
KDTree-based spatial query engine using 3D Geocentric (WGS 84) coordinates.
Handles both nearest-neighbor and radius-based lookups between satellite points and model grid nodes using a fast 3D KDTree built on ECEF (Earth-Centered, Earth-Fixed) coordinates.
__init__(model_lon, model_lat, model_height=None)
Parameters
model_lon : np.ndarray Longitudes of model mesh nodes (degrees) model_lat : np.ndarray Latitudes of model mesh nodes (degrees) model_height : np.ndarray, optional Heights of model mesh nodes above ellipsoid (meters). If None, defaults to 0 (on the ellipsoid surface).
query_nearest(sat_lon, sat_lat, sat_height, k=3)
Query for the 'k' nearest model nodes to each satellite point.
Parameters
sat_lon : np.ndarray Longitudes of satellite observations (degrees) sat_lat : np.ndarray Latitudes of satellite observations (degrees) sat_height : np.ndarray Heights/Altitudes of satellite observations (meters above ellipsoid) k : int, optional Number of nearest model neighbors (default is 3)
Returns
tuple of np.ndarray Distances (meters) and indices of nearest model nodes, shape (N, k)
query_radius(sat_lon, sat_lat, sat_height, radius_m)
Query for all model nodes within the 3D search radius.
Parameters
sat_lon : np.ndarray Longitudes of satellite observations (degrees) sat_lat : np.ndarray Latitudes of satellite observations (degrees) sat_height : np.ndarray Heights/Altitudes of satellite observations (meters above ellipsoid) radius_m : float Search radius in meters
Returns
tuple of list[np.ndarray], list[np.ndarray] - Distances (in meters) to all matched model nodes per point - Corresponding indices of model nodes per point
inverse_distance_weights(distances, power=1.0)
Compute inverse distance weights (IDW) with configurable exponent.
Parameters
distances : np.ndarray Distance array to nearest neighbors, shape (N, k) power : float, optional Power exponent for distance weighting (default is 1.0). Use 1.0 for linear, 2.0 for quadratic, etc.
Returns
np.ndarray Normalized inverse distance weights of shape (N, k)
Notes
If a distance is exactly 0, its weight becomes 1 and all others 0.
lat_lon_to_cartesian_vec(latitude, longitude, height)
Converts geodetic coordinates (latitude, longitude, height) to geocentric Cartesian coordinates (X, Y, Z) using numpy.
This function uses the WGS 84 ellipsoid model. Assumes latitude and longitude are in degrees, height in meters. Returns (X, Y, Z) in meters.
Notes
This was suggested by Gemini as a better approximation than a perfect sphere
Temporal
ocstrack.Collocation.temporal
Functions for temporal collocation
temporal_interpolated(ds_obs, model_times, buffer, time_coord_name='time')
Perform linear time interpolation using surrounding model timestamps.
Parameters
ds_obs : xr.Dataset Observation dataset with a 'time' dimension model_times : np.ndarray Array of model time values (np.datetime64) buffer : np.timedelta64 Time buffer to include obs points near the model time window
Returns
tuple obs_sub : xr.Dataset Subset of obs data used in interpolation ib : np.ndarray Index of earlier model time ia : np.ndarray Index of later model time weights : np.ndarray Linear interpolation weights for each obs point time_deltas : np.ndarray Difference (in seconds) between obs and closest model time
temporal_nearest(ds_obs, model_times, buffer, time_coord_name='time')
Match observations observations to the nearest model timestamps.
Parameters
ds_obs : xr.Dataset Observation dataset with a 'time' dimension model_times : np.ndarray Array of model time values (np.datetime64) buffer : np.timedelta64 Time buffer to include obs points near the model time window
Returns
tuple obs_sub : xr.Dataset Subset of obs data within the buffered time range nearest_inds : np.ndarray Index of nearest model time for each obs time time_deltas : np.ndarray Difference (in seconds) between obs and matched model time