corelle.engine 3.0.0__tar.gz

corelle_engine-3.0.0/PKG-INFO

@@ -0,0 +1,31 @@
+Metadata-Version: 2.4
+Name: corelle.engine
+Version: 3.0.0
+Summary: Plate-rotation models and rotation engine for the Corelle system
+Author: Daven Quinn
+Author-email: Daven Quinn <dev@davenquinn.com>
+License-Expression: MIT
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Requires-Dist: click>=8.1,<9
+Requires-Dist: corelle-client>=2.2,<3
+Requires-Dist: macrostrat-database>=4.7,<5
+Requires-Dist: macrostrat-utils>=1.0.1,<2
+Requires-Dist: numba>=0.63.1,<0.69
+Requires-Dist: numpy>=2.0.0,<3
+Requires-Dist: numpy-quaternion>=2024.0.13,<2025.0.0
+Requires-Dist: rich>12.6.0,<16
+Requires-Dist: typer>0.12.0,<0.30
+Requires-Dist: geoalchemy2>0.14.0,<0.22
+Requires-Dist: sqlalchemy>=2.0.25,<2.1
+Requires-Dist: fiona>=1.10.1,<2 ; extra == 'import'
+Requires-Dist: pyyaml>=6.0,<7 ; extra == 'import'
+Requires-Dist: wget>=3.2,<4 ; extra == 'import'
+Requires-Dist: ipython>=8.0,<10 ; extra == 'shell'
+Requires-Python: >=3.10, <4
+Provides-Extra: import
+Provides-Extra: shell

corelle_engine-3.0.0/corelle/engine/__init__.py

@@ -0,0 +1 @@
+from .cli import cli

corelle_engine-3.0.0/corelle/engine/cache.py

@@ -0,0 +1,78 @@
+"""Precomputed rotations, stored in `corelle.rotation_cache`."""
+
+import numpy as N
+from sqlalchemy.sql import text
+
+from .query import get_sql
+from .rotate import get_rotation_series
+
+update_derived = get_sql("update-cache")
+
+_cached_rotation = text(
+    """
+    SELECT c.rotation
+    FROM corelle.rotation_cache c
+    JOIN corelle.model m ON m.id = c.model_id
+    WHERE m.name = :model_name AND c.plate_id = :plate_id AND c.t_step = :time
+    """
+)
+
+_insert_rotation = text(
+    """
+    INSERT INTO corelle.rotation_cache (model_id, plate_id, t_step, rotation)
+    VALUES (:model_id, :plate_id, :t_step, :rotation)
+    """
+)
+
+
+def get_from_cache(conn, cache_args):
+    (model_name, plate_id, time) = cache_args
+    params = dict(model_name=model_name, plate_id=plate_id, time=time)
+    res = conn.execute(_cached_rotation, params).scalar()
+    if res is not None:
+        return N.quaternion(*res)
+    return None
+
+
+def build_rotation_caches(conn):
+    conn.execute(text("DELETE FROM corelle.rotation_cache"))
+    models = conn.execute(text("SELECT * FROM corelle.model")).fetchall()
+    for model in models:
+        build_rotation_cache(conn, model)
+
+
+def build_rotation_cache(conn, model, time_step=1):
+    min_age = model.min_age or 0
+    max_age = model.max_age or 1000
+
+    # Get model steps every 1 Myr
+    t_steps = list(range(int(min_age), int(max_age) + 1, time_step))
+
+    conn.execute(
+        text("DELETE FROM corelle.rotation_cache WHERE model_id = :model_id"),
+        dict(model_id=model.id),
+    )
+
+    for tstep in get_rotation_series(conn, model.name, *t_steps):
+        time = tstep["time"]
+        print(model.name, " - ", time, " Ma")
+        rows = [
+            build_cache_row(model.id, plate_id, time, q)
+            for plate_id, q in tstep["rotations"]
+            if q is not None
+        ]
+        add_to_cache(conn, rows)
+
+
+def build_cache_row(model_id, plate_id, t_step, q):
+    return dict(
+        model_id=int(model_id),
+        plate_id=int(plate_id),
+        t_step=float(t_step),
+        rotation=[float(v) for v in (q.w, q.x, q.y, q.z)],
+    )
+
+
+def add_to_cache(conn, data):
+    if data:
+        conn.execute(_insert_rotation, data)

corelle_engine-3.0.0/corelle/engine/cli.py

@@ -0,0 +1,188 @@
+import click
+from click import group, argument, option, echo
+from macrostrat.utils import working_directory
+from contextlib import contextmanager
+from pathlib import Path
+
+
+@contextmanager
+def transaction():
+    from .database import get_database
+
+    with get_database().engine.begin() as conn:
+        yield conn
+
+
+@group()
+def cli():
+    """Register and pre-calculate plate-rotation models"""
+    pass
+
+
+@cli.command(name="init")
+@option("--drop", is_flag=True, default=False)
+def init(drop=False):
+    """Create database fixtures"""
+    from .database import initialize
+
+    initialize(drop=drop)
+
+
+file = click.Path(exists=True, dir_okay=False)
+
+
+@cli.command(name="import")
+@argument("model_name")
+@argument("plates", type=file)
+@argument("rotations", type=file)
+@option("--fields", type=file)
+@option("--overwrite", is_flag=True, default=False)
+@option("--min-age", type=float, default=None)
+@option("--max-age", type=float, default=None)
+def _import(
+    model_name,
+    plates,
+    rotations,
+    fields=None,
+    overwrite=False,
+    min_age=None,
+    max_age=None,
+):
+    """
+    Import a plate-rotation model
+    """
+    from .load_data import import_model
+
+    with transaction() as conn:
+        import_model(
+            conn,
+            model_name,
+            plates,
+            rotations,
+            fields=fields,
+            overwrite=overwrite,
+            min_age=min_age,
+            max_age=max_age,
+        )
+
+
+@cli.command(name="import-features")
+@argument("name")
+@argument("file", type=file)
+@option("--overwrite", is_flag=True, default=False)
+def _import_features(name, file, overwrite=False):
+    """
+    Import features that can be associated with the models
+    """
+    from .load_data import import_features
+
+    with transaction() as conn:
+        import_features(conn, name, file, overwrite=overwrite)
+
+
+@cli.command(name="import-starter-data")
+def import_basic():
+    """Import basic models and data"""
+    from .load_data import load_basic_data
+
+    with transaction() as conn:
+        load_basic_data(conn)
+
+
+@cli.command(name="reset-cache")
+def cache():
+    """
+    Compute and cache rotations
+    """
+    from .rotate import reset_cache
+
+    reset_cache()
+    echo("Cache was reset!")
+
+
+@cli.command(name="rotate")
+@argument("model", type=str)
+@argument("plate", type=int)
+@argument("time", type=float)
+@option("--verbose", "-v", is_flag=True, default=False)
+def rotate(model, plate, time, verbose=False):
+    """
+    Rotate a plate to a time
+    """
+    import numpy as N
+    from corelle.math import quaternion_to_euler
+    from .rotate import get_rotation
+
+    with transaction() as conn:
+        q = get_rotation(conn, model, plate, time, verbose=verbose)
+    if q is None:
+        echo(f"No rotation for plate {plate} at {time} Ma.")
+        return
+    angle = N.degrees(q.angle())
+    echo(f"Rotate {angle:.2f}° around {q.vec}")
+    lat, lon, angle = quaternion_to_euler(q)
+    echo(f"Pole: {lat:.2f}°, {lon:.2f}°; Angle: {angle:.2f}°")
+
+
+@cli.command(name="rotate-all")
+@argument("model", type=str)
+@argument("time", type=float)
+@option("--verbose", "-v", is_flag=True, default=False)
+def rotate_all(model, time, verbose=False):
+    """
+    Rotate all plates in a model to a time
+    """
+    import numpy as N
+    from .rotate import get_all_rotations
+
+    with transaction() as conn:
+        for plate_id, q in get_all_rotations(conn, model, time, verbose=verbose):
+            angle = N.degrees(q.angle())
+            echo(f"{plate_id}: rotate {angle:.2f}° around {q.vec}")
+
+
+@cli.command(name="serve")
+@option("-p", "--port", type=int, default=5000)
+@option("--debug", is_flag=True, default=False)
+def serve(**kwargs):
+    """Run the application server"""
+    from corelle.server import app
+
+    app.run(host="0.0.0.0", **kwargs)
+
+
+@cli.command(name="shell")
+def shell():
+    """Get a shell in the application context"""
+    from IPython import embed
+
+    embed()
+
+
+@cli.command(name="cache-rotations")
+def build_cache():
+    """Cache rotations for all models"""
+    from .cache import build_rotation_caches
+
+    with transaction() as conn:
+        build_rotation_caches(conn)
+
+
+def parent_dir(_start: str, name: str) -> Path:
+    start = Path(_start)
+    while start.parent is not None:
+        start = start.parent
+        if start.name == name:
+            return start
+    return None
+
+
+@cli.command(name="test", context_settings=dict(ignore_unknown_options=True))
+def test(*args):
+    """Run tests"""
+    basedir = parent_dir(__file__, "py-packages")
+
+    with working_directory(basedir):
+        from pytest import main
+
+        main(list(args))

corelle_engine-3.0.0/corelle/engine/database.py

@@ -0,0 +1,35 @@
+from functools import cache
+from os import environ
+from pathlib import Path
+
+from macrostrat.database import Database
+from macrostrat.utils import relative_path
+from sqlalchemy.exc import ProgrammingError
+
+
+def database_url():
+    return environ.get("CORELLE_DB", "postgresql:///plate-rotations")
+
+
+def get_database(url=None) -> Database:
+    """The database for the CLI and the server; engine functions take a connection."""
+    return _database(url or database_url())
+
+
+@cache
+def _database(url):
+    # Pre-ping replaces pooled connections the server has dropped.
+    return Database(url, pool_pre_ping=True)
+
+
+def initialize(db=None, drop=False):
+    db = db or get_database()
+    dn = Path(relative_path(__file__, "schema"))
+    if drop:
+        try:
+            db.run_sql("DROP SCHEMA corelle CASCADE")
+        except ProgrammingError:
+            pass
+
+    for file in sorted(dn.glob("*.sql")):
+        db.run_sql(file, has_server_binds=False)

corelle_engine-3.0.0/corelle/engine/load_data.py

@@ -0,0 +1,350 @@
+from sqlalchemy import func, text
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.dialects.postgresql import insert
+from time import perf_counter
+from json import dumps
+import fiona
+from click import echo, style
+import json
+import yaml
+from os.path import splitext
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from macrostrat.utils import working_directory
+from wget import download
+from contextlib import contextmanager
+from geoalchemy2 import functions as gfunc
+from geoalchemy2.shape import from_shape
+from shapely.geometry import shape, MultiPolygon
+
+from .query import get_sql
+from .storage import reflect_table
+
+
+def create_model(conn, name, **kwargs):
+    model = reflect_table(conn, "model")
+    try:
+        # A savepoint, so a conflict doesn't abort the caller's transaction.
+        with conn.begin_nested():
+            conn.execute(
+                insert(model)
+                .values(name=name, **kwargs)
+                .on_conflict_do_update(index_elements=(model.c.id,), set_=kwargs)
+            )
+    except IntegrityError:
+        pass
+    return conn.execute(model.select().where(model.c.name == name)).first()[0]
+
+
+def pg_geometry(feature):
+    """Convert a GeoJSON feature to a PostGIS geometry"""
+    geom = shape(feature.geometry)
+    # ensure the geometry is a valid multipolygon
+    return gfunc.ST_Multi(gfunc.ST_MakeValid(from_shape(geom, srid=4326)))
+
+
+def insert_rows(conn, statement, rows, chunk_size=1000):
+    """Insert in chunks: one VALUES list overflows psycopg's 65,535 parameters.
+
+    Not executemany, since geometries are SQL expressions rather than parameters.
+    """
+    rows = list(rows)
+    keys = set().union(*rows) if rows else set()
+    rows = [{k: row.get(k) for k in keys} for row in rows]
+    for i in range(0, len(rows), chunk_size):
+        conn.execute(statement.values(rows[i : i + chunk_size]))
+
+
+def import_plate(model_id, feature, fields=None):
+    def field(field_id):
+        # If we have provided a mapping of fields for the input
+        # geojson file, we should use these fields for
+        if fields and field_id in fields:
+            field_id = fields[field_id]
+        return feature["properties"].get(field_id, None)
+
+    plate_id = field("id") or feature.get("id", None)
+    if plate_id is None:
+        raise Exception(f"Invalid plate ID: {feature}")
+
+    plate = dict(
+        id=plate_id,
+        model_id=model_id,
+        parent_id=field("parent_id"),
+        name=field("name"),
+        cotid=field("cotid"),
+        coid=field("coid"),
+    )
+
+    young_lim = field("young_lim")
+    if young_lim == -999:
+        young_lim = None
+
+    plate_polygon = dict(
+        plate_id=plate_id,
+        model_id=model_id,
+        young_lim=young_lim,
+        old_lim=field("old_lim"),
+        geometry=pg_geometry(feature),
+    )
+
+    return plate, plate_polygon
+
+
+def import_plates(conn, model_id, datafile, fields={}):
+    """Import all plates for a model"""
+    if fields is None:
+        fields = {}
+
+    # A plate can have several polygons (PalaeoPlates has 390 such plates); keep
+    # one row per plate ID, the first seen.
+    plates = {}
+    plate_polygons = []
+    with fiona.open(datafile, "r") as src:
+        mrf = fields.get("mantle_reference_frame", None)
+        if mrf is not None:
+            plate = dict(id=mrf, model_id=model_id, name="Mantle reference frame")
+            plates.setdefault(mrf, plate)
+
+        # The spin axis is assumed to have plate ID 0 by default
+        spin_axis = fields.get("spin_axis", 0)
+        spin_axis_plate = dict(id=spin_axis, model_id=model_id, name="Spin axis")
+        plates.setdefault(spin_axis, spin_axis_plate)
+
+        for feature in src:
+            plate, plate_polygon = import_plate(model_id, feature, fields=fields)
+            plates.setdefault(plate["id"], plate)
+            plate_polygons.append(plate_polygon)
+
+    insert_rows(conn, insert(reflect_table(conn, "plate")), plates.values())
+    insert_rows(conn, insert(reflect_table(conn, "plate_polygon")), plate_polygons)
+
+    # For faster updates, this materialized view could become an actual table
+    conn.execute(text("REFRESH MATERIALIZED VIEW corelle.plate_polygon_cache"))
+
+
+def create_feature(dataset, feature):
+    props = feature["properties"]
+    if props is not None:
+        props = dict(props)
+
+    return dict(
+        dataset_id=dataset,
+        properties=(props or None),
+        geometry=pg_geometry(feature),
+    )
+
+
+def import_features(conn, name, features, overwrite=False):
+    start = perf_counter()
+    table = reflect_table(conn, "feature")
+    with fiona.open(features, "r") as src:
+        if overwrite:
+            conn.execute(table.delete().where(table.c.dataset_id == name))
+
+        vals = []
+        for i, feature in enumerate(src):
+            vals.append(create_feature(name, feature))
+
+    insert_rows(conn, insert(table), vals)
+
+    step1 = perf_counter()
+    elapsed = step1 - start
+    echo(
+        f"Imported {i+1} features for dataset "
+        + style(name, bold=True)
+        + f" in {elapsed:.2f} seconds"
+    )
+
+    sql = get_sql("cache-feature-dataset")
+    conn.execute(sql, dict(dataset_id=name))
+
+    elapsed = perf_counter() - step1
+    echo(f"  cached transformed features in {elapsed:.2f} seconds")
+
+
+def create_rotation_row(model_id, line):
+    """
+    Import a line from a GPlates rotations file
+    """
+    if line.startswith("*"):
+        return None, []
+
+    data, meta = line.strip().split("!", 1)
+    row = data.split()
+
+    if int(row[0]) == 999:
+        return None, []
+
+    # Insert first plate id
+    plate_id = row[0]
+    ref_plate_id = row[5]
+    # Not sure that every plate_id-ref_plate_id
+    # pair is actually unique as would be needed
+    # here...
+    ref_plate = dict(id=ref_plate_id, model_id=model_id)
+    plate = dict(id=plate_id, model_id=model_id, parent_id=ref_plate_id)
+
+    rotation = dict(
+        plate_id=plate_id,
+        model_id=model_id,
+        t_step=row[1],
+        latitude=row[2],
+        longitude=row[3],
+        angle=row[4],
+        ref_plate_id=ref_plate_id,
+        # Postgres chokes on NULL characters in strings
+        metadata=meta.replace("\x00", ""),
+    )
+
+    # Return the rotation and the plates that it refers to
+    return rotation, [ref_plate, plate]
+
+
+def import_rotations(conn, model_id, filename):
+    """
+    Import a set of plate rotations defined
+    in a GPlates rotation file structure.
+
+    Constructed on 2019-03-15 for use with
+    PaleoPlates data, so it might not work
+    well for other `.rot` files in its current
+    form.
+    """
+    rotations = []
+    plates = []
+    start = perf_counter()
+    with open(filename, "r") as f:
+        for i, line in enumerate(f):
+            rotation, next_plates = create_rotation_row(model_id, line)
+            if rotation is not None:
+                rotations.append(rotation)
+            plates.extend(next_plates)
+
+    plate = reflect_table(conn, "plate")
+    insert_rows(
+        conn,
+        insert(plate).on_conflict_do_nothing(
+            index_elements=(plate.c.id, plate.c.model_id)
+        ),
+        plates,
+    )
+    rotation = reflect_table(conn, "rotation")
+    insert_rows(conn, insert(rotation).on_conflict_do_nothing(), rotations)
+
+    elapsed = perf_counter() - start
+    print(f"Imported {i+1} rotations in {elapsed:.2f} seconds")
+
+
+def import_model(
+    conn,
+    name,
+    plates,
+    rotations,
+    fields=None,
+    overwrite=False,
+    min_age=None,
+    max_age=None,
+):
+    print("Importing model", name)
+    q = text("SELECT count(*) FROM corelle.model WHERE name=:name")
+    res = conn.execute(q, dict(name=name)).scalar()
+    if res == 1 and not overwrite:
+        print("Model has already been imported.")
+        return
+
+    model_id = create_model(conn, name, min_age=min_age, max_age=max_age)
+    import_plates(conn, model_id, plates, fields=load_fields(fields))
+    import_rotations(conn, model_id, rotations)
+    # Right now we don't cache rotations on model import, but we could.
+
+
+def load_fields(fn):
+    if not fn:
+        return None
+    ext = splitext(fn)[1]
+    with open(fn, "r") as f:
+        if ext == ".json":
+            return json.load(f)
+        if ext in [".yaml", ".yml"]:
+            return yaml.load(f, Loader=yaml.SafeLoader)
+    return None
+
+
+corelle_data_dir = Path(__file__).parent.parent.parent.parent.parent / "data"
+
+
+def load_basic_data(conn):
+    """Load basic model and feature datasets"""
+    with working_directory(str(corelle_data_dir)):
+        import_model(
+            conn,
+            "PalaeoPlates",
+            "eglington/PlatePolygons2016All.json",
+            "eglington/T_Rot_Model_PalaeoPlates_2019_20190302_experiment.rot",
+            fields="eglington-fields.yaml",
+        )
+
+        import_model(
+            conn,
+            "Seton2012",
+            "seton_2012.geojson",
+            "Seton_etal_ESR2012_2012.1.rot",
+            fields="seton-fields.yaml",
+            min_age=0,
+            max_age=200,
+        )
+
+        import_model(
+            conn,
+            "Wright2013",
+            "wright_plates.geojson",
+            "wright_2013.rot",
+            fields="wright-fields.yaml",
+            min_age=0,
+            max_age=550,
+        )
+
+        import_model(
+            conn,
+            "Scotese",
+            "scotese.geojson",
+            "scotese.rot",
+            fields="scotese-fields.yaml",
+            min_age=0,
+            max_age=550,
+        )
+
+    # Import Natural Earth features
+    # Download natural earth data to temp file
+    prefix = "https://raw.githubusercontent.com/martynafford/natural-earth-geojson/master/110m/"
+
+    # Create a temp directory
+    with download_context():
+        for dataset in [
+            prefix + "cultural/ne_110m_populated_places_simple.json",
+            prefix + "physical/ne_110m_land.json",
+        ]:
+            fn = dataset.split("/")[-1]
+            dsn = fn.split(".")[0]
+
+            # Check if this dataset has already been imported
+            q = text(
+                "SELECT count(*) FROM (SELECT DISTINCT dataset_id FROM corelle.feature WHERE dataset_id=:name) AS a"
+            )
+            res = conn.execute(q, dict(name=dsn)).scalar()
+            if res == 1:
+                continue
+
+            # Download the dataset
+            download(dataset, out=fn)
+
+            # Import the dataset
+            import_features(conn, dsn, fn, overwrite=True)
+
+
+@contextmanager
+def download_context():
+    with TemporaryDirectory() as tmpdir:
+        with working_directory(tmpdir):
+            yield tmpdir

corelle_engine-3.0.0/corelle/engine/query/__init__.py

@@ -0,0 +1,7 @@
+from macrostrat.utils import relative_path
+from sqlalchemy import text
+
+
+def get_sql(key):
+    fn = relative_path(__file__, key + ".sql")
+    return text(open(fn).read())

corelle_engine-3.0.0/corelle/engine/query/active-plates-at-time.sql

@@ -0,0 +1,11 @@
+SELECT
+  p.id
+FROM corelle.plate p
+JOIN corelle.plate_polygon pp
+  ON p.id = pp.plate_id
+ AND p.model_id = pp.model_id
+JOIN corelle.model m
+  ON p.model_id = m.id
+WHERE m.name = :model_name
+  AND coalesce(pp.old_lim, m.max_age) > :time
+  AND coalesce(pp.young_lim, m.min_age) < :time