Source code for okada.model

"""

:copyright:
    2024, Conor A. Bacon.
:license:
    GNU General Public License, Version 3
    (https://www.gnu.org/licenses/gpl-3.0.html)

"""

import itertools
from dataclasses import dataclass
import pathlib
import tomllib

import numpy as np
import pandas as pd
import pyproj
from pyproj.enums import TransformDirection

from okada.elements import coulomb2okadapy


[docs] @dataclass class Model: """Class to encapsulate model information for `OkadaPy`.""" poisson_ratio: float youngs_modulus: float friction_coefficient: float elements: list x_coords: np.ndarray y_coords: np.ndarray size: list[float] | None = None xsection: list[float] | None = None map_: list[float] | None = None transformer: pyproj.Transformer | None = None @property def grid_bounds_xy(self) -> list[tuple[float], tuple[float]]: """Get x and y coordinate bounds for the model grid.""" bottom_left_corner = (min(self.x_coords), min(self.y_coords)) top_right_corner = (max(self.x_coords), max(self.y_coords)) return [bottom_left_corner, top_right_corner] @property def grid_bounds_coords(self) -> list[tuple[float], tuple[float]]: """Get lat/lon coordinate bounds for the model grid.""" (min_lon, min_lat), (max_lon, max_lat) = [ self.transformer.transform( x * 1000, y * 1000, direction=TransformDirection.INVERSE ) for x, y in self.grid_bounds_xy ] return [min_lon, max_lon, min_lat, max_lat] @property def grid_xy(self) -> tuple[np.ndarray, np.ndarray]: """Get x and y coordinates of all nodes in model grid.""" X, Y = np.mgrid[ min(self.x_coords) : max(self.x_coords) : len(set(self.x_coords)) * 1j, min(self.y_coords) : max(self.y_coords) : len(set(self.y_coords)) * 1j, ] return X, Y @property def grid_coords(self) -> tuple[np.ndarray, np.ndarray]: """Get lat/lon coordinates of all nodes in model grid.""" if self.transformer is None: print("No coordinate transformation defined for this model.") return grid_x, grid_y = self.grid_xy grid_lons, grid_lats = self.transformer.transform( grid_x * 1000, grid_y * 1000, direction=TransformDirection.INVERSE, ) return grid_lons, grid_lats @property def raw_elements(self) -> np.ndarray: """ Prepare all model elements for computation, which assumes parameters have specific positions in a flattened, contiguous array. """ return np.asarray([element.raw_input for element in self.elements]).flatten()
[docs] def read(model_file: str, file_format: str = None) -> Model: """ Parse a model file into the Model dataclass. Arguments --------- model_file: Path to a file containing the model to load. file_format: A string identifier for a specific file format e.g. "COULOMB" Returns ------- model: The model in the `OkadaPy` Model format. Raises ------ FileNotFoundError: If the specified model_file does not exist. Exception: If the model file format could not be detected. """ print(f"\tReading in model from {model_file}...") model_file = pathlib.Path(model_file) if not model_file.is_file(): raise FileNotFoundError # Read file. Will try to automatically detect the file format if model_file.suffix == ".toml" or file_format == "TOML": print("\t\t.toml file identified.") model = _read_toml(model_file) elif model_file.suffix == ".inp" or file_format == "COULOMB": print("\t\tCoulomb-style .inp file identified.") model = _read_coulomb(model_file) else: raise Exception("Could not detect file format—please specify.") print("\t...read successful.") return model
def _read_coulomb(model_file: pathlib.Path) -> Model: """ Parse the model and model parameters from a file that conforms to the Coulomb3 model file format. Arguments --------- model_file: Path to a Coulomb3 model file containing the model to load. Returns ------- model: The model in the `OkadaPy` Model format. """ with open(model_file, "r") as f: lines = f.readlines() header, body, *tail = [ list(group) for k, group in itertools.groupby(lines, lambda x: x == "\n") if not k ] tail = [section for sublist in tail for section in sublist] # Parse information from model file header poisson_ratio = float(header[3].strip().split()[1]) youngs_modulus = float(header[4].strip().split()[1]) friction_coefficient = float(header[6].strip().split()[1]) # Parse model elements from model file body raw_elements = ( pd.DataFrame([line.strip().split() for line in body[2:]]) .drop(columns=0) .values.astype(np.float64) ) elements = [coulomb2okadapy(element) for element in raw_elements] # Parse grid information from model file tail grid, size, xsection, map_ = tail[:7], tail[7:11], tail[11:19], tail[19:] grid_keys = ["x_min", "y_min", "x_max", "y_max", "x_inc", "y_inc"] grid = { key: float(model_part.split("=")[1].strip()) for key, model_part in zip(grid_keys, grid[1:]) } x_coords = np.arange(grid["x_min"], grid["x_max"] + grid["x_inc"], grid["x_inc"]) x_coords = np.ascontiguousarray(x_coords, dtype=np.float64) y_coords = np.arange(grid["y_min"], grid["y_max"] + grid["y_inc"], grid["y_inc"]) y_coords = np.ascontiguousarray(y_coords, dtype=np.float64) x_coords, y_coords = np.meshgrid(x_coords, y_coords, indexing="ij") map_keys = ["lon_min", "lon_max", "lon_zero", "lat_min", "lat_max", "lat_zero"] map_ = { key: float(map_part.split("=")[1].strip()) for key, map_part in zip(map_keys, map_[1:]) } model = Model( poisson_ratio, youngs_modulus, friction_coefficient, elements, x_coords.flatten(), y_coords.flatten(), size, xsection, map_, ) return model def _read_toml(model_file: pathlib.Path) -> Model: """ Parse the model and model parameters from a .toml file (defined by `OkadaPy`). Arguments --------- model_file: Path to an OkadaPy .toml model file containing the model to load. Returns ------- model: The model in the `OkadaPy` Model format. """ with model_file.open("rb") as f: tf = tomllib.load(f) elements = pd.read_csv(model_file.parent / tf["model"]["elements_file"]) elements = elements.values.flatten() elements = elements.astype(np.float64) if "grid" in tf.keys(): x_coords = np.arange(*tf["grid"]["x_range"]) y_coords = np.arange(*tf["grid"]["y_range"]) z_coords = np.arange(*tf["grid"]["z_range"]) elif "coords" in tf.keys(): x_coords = tf["coords"]["x"] y_coords = tf["coords"]["y"] z_coords = tf["coords"]["z"] x_coords = np.ascontiguousarray(x_coords, dtype=np.float64) y_coords = np.ascontiguousarray(y_coords, dtype=np.float64) z_coords = np.ascontiguousarray(z_coords, dtype=np.float64) x_coords, y_coords = np.meshgrid(x_coords, y_coords, indexing="ij") model = Model( tf["model"]["poisson_ratio"], float(tf["model"]["youngs_modulus"]), tf["model"]["friction_coefficient"], elements, x_coords.flatten(), y_coords.flatten(), ) return model