diff --git a/imdb/imdb.py b/imdb/imdb.py index 1e15d86..c828763 100644 --- a/imdb/imdb.py +++ b/imdb/imdb.py @@ -236,6 +236,10 @@ def _id_map(self, table: str, id_col: str, int_col: str) -> pd.Series: df = self.con.table(table).select(id_col, int_col).to_pandas() return df.set_index(id_col)[int_col] + def _columns(self, table: str) -> list[str]: + """The columns `table` has in this database (older schema versions lack some).""" + return list(self.con.table(table).columns) + def add_events(self, df: pd.DataFrame) -> None: """Insert new events. @@ -506,6 +510,190 @@ def delete_event(self, event_id: str) -> None: for statement in statements: self.con.raw_sql(statement, parameters=[event_id]) + def _grid(self, table: str, index_col: str, value_col: str) -> np.ndarray: + """Return the period or frequency grid, in index order.""" + df = self.con.table(table).to_pandas().sort_values(index_col) + return df[value_col].to_numpy(dtype=float) + + def _ims_by_record( + self, table: str, column: str, record_int_ids: np.ndarray, width: int + ) -> tuple[np.ndarray | None, np.ndarray, np.ndarray]: + """Return one array IM table (and its sigma) as matrices aligned to `record_int_ids`. + + Parameters + ---------- + table : str + `"psa_ims"` or `"fas_ims"`. + column : str + `"pSA"` or `"FAS"`. + record_int_ids : np.ndarray + Records to return rows for, in order. + width : int + Grid length (number of periods or frequencies). + + Returns + ------- + tuple of np.ndarray + `(values, sigma, has_sigma)`: `values` and `sigma` shaped + `(len(record_int_ids), width)` with NaN rows for records missing from the + table (`values` is None when no record has a row), and `has_sigma` a bool per + record that has a sigma array. + """ + values = np.full((len(record_int_ids), width), np.nan) + sigma = np.full((len(record_int_ids), width), np.nan) + has_sigma = np.zeros(len(record_int_ids), dtype=bool) + rows = ( + self.con.table(table) + .filter(_.record_int_id.isin(record_int_ids.tolist())) + .to_pandas() + .set_index("record_int_id") + ) + if rows.empty: + return None, sigma, has_sigma + position = pd.Series(np.arange(len(record_int_ids)), index=record_int_ids) + positions = position.loc[rows.index].to_numpy() + for ix, row_values, row_sigma in zip( + positions, rows[column], rows[f"{column}_sigma"], strict=True + ): + values[ix] = np.asarray(row_values, dtype=float) + if isinstance(row_sigma, list | np.ndarray) and len(row_sigma): + sigma[ix] = np.asarray(row_sigma, dtype=float) + has_sigma[ix] = True + return values, sigma, has_sigma + + def merge_from(self, other: "IMDB", replace: bool = True) -> list[str]: + """Copy every event of another IMDB into this one, with everything derived from it. + + Copies the events, their realisations, site-event distances, records and IMs. + Sites are matched by `site_id`: new sites are added, and existing ones take + `other`'s values where `other` has them. A NULL in `other` keeps this database's + value, since a shared site is referenced by other events here too. Both databases + must share the same period (and frequency) grids, and `other`'s components must + all be in this database. + + Not atomic as a whole (`add_records` commits its own transaction), so merge into a + copy of the target and swap it in once `validate` passes. + + Parameters + ---------- + other : IMDB + An open database to copy from (read-only is fine). + replace : bool + Delete events that already exist here before copying them, by default True. + When False, an event that already exists is an error. + + Returns + ------- + list of str + The merged event ids. + + Raises + ------ + ValueError + If the grids or components don't match, or an event already exists here and + `replace` is False. + """ + # The IM arrays are indexed by the grids, so they must match exactly + periods = other._grid("periods", "period_index", "period") + if not np.array_equal(periods, self._grid("periods", "period_index", "period")): + raise ValueError( + "period grids differ, can't merge pSA between these databases" + ) + frequencies = other._grid("frequencies", "freq_index", "frequency") + if len(frequencies) and not np.array_equal( + frequencies, self._grid("frequencies", "freq_index", "frequency") + ): + raise ValueError( + "frequency grids differ, can't merge FAS between these databases" + ) + unknown = set(other.db_meta["components"].split(",")) - set( + self.db_meta["components"].split(",") + ) + if unknown: + raise ValueError(f"components not in this database: {sorted(unknown)}") + + events = other.get_events().reset_index() + event_ids = events["event_id"].tolist() + existing = sorted(set(event_ids) & set(self.get_events().index)) + if existing and not replace: + raise ValueError(f"events already in this database: {existing}") + for event_id in existing: + self.delete_event(event_id) + + # Sites: add new ones, update existing ones in place without letting a missing + # incoming value (NULL) erase a known one + sites = other.get_sites().reset_index().drop(columns="site_int_id") + sites = sites[[c for c in sites.columns if c in self._columns("sites")]] + known = sites["site_id"].isin(self.get_sites().index) + if (~known).any(): + self.add_sites(sites[~known]) + if known.any(): + # The underlying duckdb connection, to UPDATE ... FROM a DataFrame + raw = self.con.con + raw.register("_merge_sites", sites[known]) + set_clause = ", ".join( + f"{c} = COALESCE(m.{c}, sites.{c})" + for c in sites.columns + if c != "site_id" + ) + try: + raw.execute( + f"UPDATE sites SET {set_clause} FROM _merge_sites m WHERE sites.site_id = m.site_id" + ) + finally: + raw.unregister("_merge_sites") + + self.add_events(events.drop(columns="event_int_id")) + event_id_by_int = events.set_index("event_int_id")["event_id"] + realisations = other.get_realisations().reset_index() + realisations["event_id"] = event_id_by_int.loc[ + realisations["event_int_id"] + ].to_numpy() + self.add_realisations(realisations.drop(columns=["rel_int_id", "event_int_id"])) + + site_event = other.get_site_event(event_ids=event_ids) + if not site_event.empty: + self.add_site_event(site_event) + + records = other.get_records(event_ids=event_ids) + if records.empty: + return event_ids + record_int_ids = records.index.to_numpy() + scalars = other.get_scalars(sigma=True, event_ids=event_ids) + has_scalars = np.isin(record_int_ids, scalars.index.to_numpy()) + records = records.join(scalars, how="left") + scalar_cols = list(scalars.columns) + psa, psa_sigma, has_psa_sigma = other._ims_by_record( + "psa_ims", "pSA", record_int_ids, len(periods) + ) + fas, fas_sigma, has_fas_sigma = other._ims_by_record( + "fas_ims", "FAS", record_int_ids, len(frequencies) + ) + + # Insert records in groups that share which IM rows they have, so each copy + # matches its source exactly: no empty scalars rows, no sigma on simulated records + groups = np.stack([has_scalars, has_psa_sigma, has_fas_sigma], axis=1) + for key in np.unique(groups, axis=0): + with_scalars, with_psa_sigma, with_fas_sigma = (bool(v) for v in key) + ix = np.flatnonzero((groups == key).all(axis=1)) + df = records.iloc[ix].reset_index(drop=True).drop(columns="event_id") + if not with_scalars: + df = df.drop(columns=scalar_cols) + else: + sigma_cols = [c for c in scalar_cols if c.endswith("_sigma")] + df = df.drop(columns=[c for c in sigma_cols if df[c].isna().all()]) + self.add_records( + df, + pSA=None if psa is None else psa[ix], + FAS=None if fas is None else fas[ix], + pSA_sigma=psa_sigma[ix] if with_psa_sigma and psa is not None else None, + FAS_sigma=fas_sigma[ix] if with_fas_sigma and fas is not None else None, + ) + logger.info( + "merged %d event(s) from %s into %s", len(event_ids), other.path, self.path + ) + return event_ids + def validate(self) -> list[str]: """Check the invariants the schema itself cannot enforce: unique ids, valid FKs. diff --git a/pyproject.toml b/pyproject.toml index 956b333..6fba391 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -101,4 +101,6 @@ exclude = [ '\.__repr__$', '\.__exit__$', '\.__enter__$', + # private helpers only need a one-line docstring + '\._[A-Za-z]\w*$', ] diff --git a/tests/test_merge.py b/tests/test_merge.py new file mode 100644 index 0000000..76c4280 --- /dev/null +++ b/tests/test_merge.py @@ -0,0 +1,244 @@ +"""Tests for `IMDB.merge_from`: copying events between databases.""" + +import numpy as np +import pandas as pd +import pytest + +from imdb import IMDB +from tests.conftest import COMPONENTS, FREQUENCIES, PERIODS + + +def _single_event_db( + path, event_id="eventC", site_lat=-43.9, pga=0.25, extra_site=True +): + """A database holding one simulated event at siteA (shared) and siteD (new).""" + db = IMDB.create( + path, periods=PERIODS, frequencies=FREQUENCIES, components=COMPONENTS + ) + db.add_events( + pd.DataFrame( + { + "event_id": [event_id], + "magnitude": [5.5], + "domain_wkt": ["POLYGON ((0 0, 1 0, 1 1, 0 0))"], + } + ) + ) + db.add_realisations( + pd.DataFrame({"rel_id": [event_id], "event_id": [event_id], "rake": [90.0]}) + ) + site_ids = ["siteA", "siteD"] if extra_site else ["siteA"] + db.add_sites( + pd.DataFrame( + { + "site_id": site_ids, + "lat": [site_lat, -44.0][: len(site_ids)], + "lon": [172.0, 172.1][: len(site_ids)], + "vs30": [400.0, 250.0][: len(site_ids)], + } + ) + ) + db.add_site_event( + pd.DataFrame( + { + "site_id": site_ids, + "event_id": [event_id] * len(site_ids), + "rrup": np.arange(len(site_ids), dtype=float) + 5, + } + ) + ) + rows = [(site_id, component) for site_id in site_ids for component in COMPONENTS] + db.add_records( + pd.DataFrame( + { + "rel_id": event_id, + "site_id": [r[0] for r in rows], + "component": [r[1] for r in rows], + "kind": "simulated", + "PGA": [pga] * len(rows), + } + ), + pSA=np.full((len(rows), len(PERIODS)), pga), + FAS=np.full((len(rows), len(FREQUENCIES)), pga / 10), + ) + return db + + +def test_merge_adds_event_and_sites(db, tmp_path): + """A new event, its new site and all its IMs are copied; shared sites are updated.""" + with _single_event_db(tmp_path / "src.duckdb") as src: + merged = db.merge_from(src) + + assert merged == ["eventC"] + assert db.validate() == [] + assert set(db.get_events().index) == {"eventA", "eventB", "eventC"} + sites = db.get_sites() + assert "siteD" in sites.index + assert sites.loc["siteA", "lat"] == pytest.approx(-43.9) + assert sites.loc["siteA", "vs30"] == pytest.approx(400.0) + + records = db.get_records(event_ids=["eventC"]) + assert len(records) == 4 + scalars = db.get_scalars(event_ids=["eventC"]) + assert np.allclose(scalars["PGA"], 0.25) + psa = db.get_psa(event_ids=["eventC"]) + assert np.allclose(psa.to_numpy(), 0.25) + fas = db.get_fas(event_ids=["eventC"]) + assert np.allclose(fas.to_numpy(), 0.025) + assert db.get_site_event(event_ids=["eventC"])["rrup"].tolist() == [5.0, 6.0] + assert db.get_realisations().loc["eventC", "rake"] == pytest.approx(90.0) + assert db.get_events().loc["eventC", "domain_wkt"].startswith("POLYGON") + + +def test_merge_keeps_site_values_missing_from_incoming(db, tmp_path): + """A NULL site value in the incoming file doesn't erase a known value here.""" + with _single_event_db(tmp_path / "rich.duckdb") as src: + db.merge_from(src) + assert db.get_sites().loc["siteA", "vs30"] == pytest.approx(400.0) + + # A later event that references siteA without a site characterisation + with _single_event_db( + tmp_path / "sparse.duckdb", event_id="eventE", site_lat=-43.8 + ) as src: + src.con.raw_sql("UPDATE sites SET vs30 = NULL") + db.merge_from(src) + + sites = db.get_sites() + assert sites.loc["siteA", "vs30"] == pytest.approx(400.0) + # Values the incoming file does have still win + assert sites.loc["siteA", "lat"] == pytest.approx(-43.8) + assert db.validate() == [] + + +def test_merge_replaces_existing_event(db, tmp_path): + """Merging an event again replaces it rather than duplicating records.""" + with _single_event_db(tmp_path / "first.duckdb", pga=0.1) as src: + db.merge_from(src) + with _single_event_db(tmp_path / "second.duckdb", pga=0.5) as src: + db.merge_from(src) + + assert db.validate() == [] + assert len(db.get_records(event_ids=["eventC"])) == 4 + assert np.allclose(db.get_scalars(event_ids=["eventC"])["PGA"], 0.5) + # Other events untouched + assert len(db.get_records(event_ids=["eventA"])) == 12 + + +def test_merge_without_replace_rejects_existing_event(db, tmp_path): + """replace=False refuses to overwrite an event that already exists.""" + with _single_event_db(tmp_path / "first.duckdb") as src: + db.merge_from(src) + with ( + _single_event_db(tmp_path / "second.duckdb") as src, + pytest.raises(ValueError, match="already in this database"), + ): + db.merge_from(src, replace=False) + + +def test_merge_rejects_different_period_grid(db, tmp_path): + """pSA arrays can't be merged across different period grids.""" + src = IMDB.create( + tmp_path / "src.duckdb", periods=[0.1, 1.0], components=COMPONENTS + ) + with src, pytest.raises(ValueError, match="period grids differ"): + db.merge_from(src) + + +def test_merge_rejects_unknown_components(db, tmp_path): + """Components the target doesn't hold can't be merged in.""" + src = IMDB.create( + tmp_path / "src.duckdb", + periods=PERIODS, + frequencies=FREQUENCIES, + components=("rotd50",), + ) + with src, pytest.raises(ValueError, match="components not in this database"): + db.merge_from(src) + + +def test_merge_into_empty_database(tmp_path): + """Merging into a freshly created database with the same grids works.""" + target = IMDB.create( + tmp_path / "target.duckdb", + periods=PERIODS, + frequencies=FREQUENCIES, + components=COMPONENTS, + ) + with target, _single_event_db(tmp_path / "src.duckdb", extra_site=False) as src: + assert target.merge_from(src) == ["eventC"] + assert target.validate() == [] + assert list(target.get_sites().index) == ["siteA"] + + +def test_merge_copies_mixed_kinds_exactly(tmp_path): + """GMM records keep their sigma, simulated ones stay sigma-free, and records without + scalars don't gain empty scalar rows.""" + src = IMDB.create( + tmp_path / "src.duckdb", + periods=PERIODS, + frequencies=FREQUENCIES, + components=COMPONENTS, + ) + src.add_events(pd.DataFrame({"event_id": ["eventM"]})) + src.add_realisations(pd.DataFrame({"rel_id": ["eventM"], "event_id": ["eventM"]})) + src.add_sites(pd.DataFrame({"site_id": ["s1"], "lat": [-43.0], "lon": [172.0]})) + # Simulated record with scalars + pSA, and a pSA-only simulated record + src.add_records( + pd.DataFrame( + { + "rel_id": ["eventM"], + "site_id": ["s1"], + "component": ["000"], + "kind": "simulated", + "PGA": [0.1], + } + ), + pSA=np.full((1, len(PERIODS)), 0.1), + ) + src.add_records( + pd.DataFrame( + { + "rel_id": ["eventM"], + "site_id": ["s1"], + "component": ["090"], + "kind": "simulated", + } + ), + pSA=np.full((1, len(PERIODS)), 0.2), + ) + # GMM record with scalar and pSA sigma + src.add_records( + pd.DataFrame( + { + "rel_id": ["eventM"], + "site_id": ["s1"], + "component": ["000"], + "kind": "gmm", + "gmm_key": ["Bradley_2013"], + "PGA": [0.08], + "PGA_sigma": [0.6], + } + ), + pSA=np.full((1, len(PERIODS)), 0.09), + pSA_sigma=np.full((1, len(PERIODS)), 0.7), + ) + + target = IMDB.create( + tmp_path / "target.duckdb", + periods=PERIODS, + frequencies=FREQUENCIES, + components=COMPONENTS, + ) + with src, target: + target.merge_from(src) + assert target.validate() == [] + for table in ("records", "scalars_ims", "psa_ims"): + assert ( + target.con.table(table).count().to_pandas() + == src.con.table(table).count().to_pandas() + ) + gmm = target.get_psa(sigma=True, kind="gmm") + assert np.allclose(gmm[f"{PERIODS[0]}_sigma"], 0.7) + assert target.get_scalars(sigma=True, kind="gmm")["PGA_sigma"].iloc[0] == ( + pytest.approx(0.6) + )