From 96362cea455414afd9023ff1e678566e8899117d Mon Sep 17 00:00:00 2001 From: jpalm3r Date: Tue, 11 Aug 2026 10:18:07 +0200 Subject: [PATCH 1/3] Read each selected node once and share one empty frame MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two costs dominated a filtered load. A node shared by several reaches was read once per reach endpoint even though _generate_graph keeps only the first copy, and every topology-only node and gridpoint allocated its own empty DataFrame — two per reach on a large network. Cache node data by id, and hand all topology-only locations one shared frame. Boundary data stays per-reach and outside the cache. Co-Authored-By: Claude Opus 5 --- src/modelskill/model/adapters/_res1d.py | 10 ++++- src/modelskill/network.py | 26 +++++++----- tests/test_match.py | 4 +- tests/test_network.py | 54 +++++++++++++++++++++++++ 4 files changed, 82 insertions(+), 12 deletions(-) diff --git a/src/modelskill/model/adapters/_res1d.py b/src/modelskill/model/adapters/_res1d.py index 567f0d498..2c335b317 100644 --- a/src/modelskill/model/adapters/_res1d.py +++ b/src/modelskill/model/adapters/_res1d.py @@ -10,6 +10,12 @@ from modelskill.network import NetworkNode, ReachBreakPoint, NetworkReach +# Topology-only nodes and gridpoints all share this frame instead of each +# allocating its own. A large network has two per reach, which profiling showed +# to be the biggest single cost of a filtered load. Never mutate it in place. +_EMPTY_DATA = pd.DataFrame() + + def _simplify_colnames(node: ResultNode | ResultGridPoint) -> pd.DataFrame: # We remove suffixes and indexes so the columns contain only the quantity names @@ -86,7 +92,7 @@ def __init__( boundary: dict[str, pd.DataFrame] | None = None, ): self._id = id - self._data = pd.DataFrame() if data is None else data + self._data = _EMPTY_DATA if data is None else data self._boundary = {} if boundary is None else boundary @property @@ -107,7 +113,7 @@ def __init__( self, reach_id: str, chainage: float, data: pd.DataFrame | None = None ): self._id = (reach_id, chainage) - self._data = pd.DataFrame() if data is None else data + self._data = _EMPTY_DATA if data is None else data @property def id(self) -> tuple[str, float]: diff --git a/src/modelskill/network.py b/src/modelskill/network.py index 4487c5809..90bc3d848 100644 --- a/src/modelskill/network.py +++ b/src/modelskill/network.py @@ -812,21 +812,29 @@ def _load_res1d_network( # potential memory issues. For this reason, we create this intermediate step that populates # only the data in the passed nodes + # A node shared by several reaches is visited once per reach endpoint, but + # _generate_graph keeps only the first copy of its data, so read it once. + # The boundary is per-reach and must stay outside the cache. + node_data: dict[str, pd.DataFrame] = {} + def _init_node(reach: ResultReach, is_end: bool) -> Res1DNode: id = reach.end_node if is_end else reach.start_node gpt_idx = -1 if is_end else 0 if id in nodes_set: - node = res.nodes[id] - df = _simplify_colnames(node) - # Merged here rather than up front so selective loading still - # decides what is held in memory. - if extra is not None and id in extra.nodes: - df = _merge_extra_quantities( - df, _simplify_colnames(extra.nodes[id]), node_id=id - ) + if id not in node_data: + df = _simplify_colnames(res.nodes[id]) + # Merged here rather than up front so selective loading still + # decides what is held in memory. + if extra is not None and id in extra.nodes: + df = _merge_extra_quantities( + df, _simplify_colnames(extra.nodes[id]), node_id=id + ) + node_data[id] = df overlapping_gridpoint = reach.gridpoints[gpt_idx] boundary = _simplify_colnames(overlapping_gridpoint) - return Res1DNode(id, data=df, boundary={reach.name: boundary}) + return Res1DNode( + id, data=node_data[id], boundary={reach.name: boundary} + ) else: return Res1DNode(id) diff --git a/tests/test_match.py b/tests/test_match.py index 3324f4788..53dc51429 100644 --- a/tests/test_match.py +++ b/tests/test_match.py @@ -7,6 +7,7 @@ import modelskill as ms from modelskill.comparison._comparison import ItemSelection from modelskill.model.dfsu import DfsuModelResult + try: from modelskill.network import _make_basic_network except ImportError: @@ -1064,7 +1065,8 @@ def test_network_match_multi_obs_multi_model_comprehensive( def test_network_match_error_non_node_observation(network_mr, point_obs_error): """Test that non-NodeObservation raises appropriate error""" with pytest.raises( - TypeError, match="NetworkModelResult supports NodeObservation and ReachObservation" + TypeError, + match="NetworkModelResult supports NodeObservation and ReachObservation", ): ms.match(point_obs_error, network_mr) diff --git a/tests/test_network.py b/tests/test_network.py index a41fc997b..54d249d24 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -647,6 +647,60 @@ def test_from_mike_empty_nodes_and_reaches_keeps_topology_and_empty_outputs(): assert len(ds.data_vars) == 0 +@pytest.mark.skipif( + sys.version_info >= (3, 14), reason="mikeio1d requires Python < 3.14" +) +def test_from_mike_topology_only_locations_share_one_empty_frame(): + """Topology-only locations must not each allocate their own empty frame. + + That is two DataFrame allocations per reach, which profiling showed to be + the single largest cost of a filtered load (gh #679). + """ + path_to_file = "./tests/testdata/network.res1d" + + network = Network.from_mike(path_to_file, nodes=[], reaches=[]) + g = network.graph + + frames = [g.nodes[n]["data"] for n in g.nodes] + assert all(frame.empty for frame in frames) + assert len({id(frame) for frame in frames}) == 1 + + +@pytest.mark.skipif( + sys.version_info >= (3, 14), reason="mikeio1d requires Python < 3.14" +) +def test_from_mike_reads_each_selected_node_once(monkeypatch): + """A node shared by several reaches is read once, not once per reach. + + _generate_graph keeps only the first copy of a node's data, so the repeat + reads were discarded work (gh #679). The per-reach boundary reads are + genuinely distinct and must not be collapsed. + """ + from modelskill.model.adapters import _res1d + + reads: dict[str, int] = {} + unpatched = _res1d._simplify_colnames + + def counting_simplify_colnames(location, *args, **kwargs): + name = type(location).__name__ + reads[name] = reads.get(name, 0) + 1 + return unpatched(location, *args, **kwargs) + + monkeypatch.setattr(_res1d, "_simplify_colnames", counting_simplify_colnames) + + path_to_file = "./tests/testdata/network.res1d" + network = Network.from_mike(path_to_file, reaches=[]) + + n_unique_nodes = len( + {reach.start.id for reach in network._reaches.values()} + | {reach.end.id for reach in network._reaches.values()} + ) + n_reach_endpoints = 2 * len(network._reaches) + + assert reads["ResultNode"] == n_unique_nodes + assert reads["ResultGridPoint"] == n_reach_endpoints + + # --------------------------------------------------------------------------- # Optional reach length # --------------------------------------------------------------------------- From e88e873e8e58ca99837a51c9c026896528c56f85 Mon Sep 17 00:00:00 2001 From: jpalm3r Date: Tue, 11 Aug 2026 10:19:29 +0200 Subject: [PATCH 2/3] Add a quantities filter to from_mike and from_epanet A calibration loop that rebuilds the network per trial but scores only one quantity still paid to read them all. Accept a quantities argument on both constructors and thread it down to the per-location read. A location that carries none of the requested quantities becomes topology-only rather than an error, so this composes with nodes and reaches on files where nodes and reaches hold different quantities. Reading every quantity stays a single interop call. Co-Authored-By: Claude Opus 5 --- docs/user-guide/network.qmd | 10 ++- src/modelskill/model/adapters/_res1d.py | 28 ++++++- src/modelskill/network.py | 50 ++++++++++++- tests/test_network.py | 98 +++++++++++++++++++++++++ 4 files changed, 177 insertions(+), 9 deletions(-) diff --git a/docs/user-guide/network.qmd b/docs/user-guide/network.qmd index a769d9c16..e769cc28a 100644 --- a/docs/user-guide/network.qmd +++ b/docs/user-guide/network.qmd @@ -250,12 +250,13 @@ A MIKE 1D network contains multiple levels that are unified into a generic netwo Large result files can contain thousands of nodes and gridpoints. Loading all of that data into memory is slow and may cause memory issues — especially when you only need the timeseries at a handful of nodes where observations exist. -Both constructors accept the same two optional arguments to restrict what gets loaded: +Both constructors accept the same three optional arguments to restrict what gets loaded: | Argument | Type | Effect | |---|---|---| | `nodes` | `None` \| `str` \| `list[str]` | Control which nodes have timeseries data loaded. `None` (default) loads all nodes; `[]` skips all node data; a name or list loads only those nodes. | | `reaches` | `None` \| `str` \| `list[str]` | Control which reaches have intermediate gridpoint data populated. `None` (default) loads everything; `[]` skips all gridpoints; a name or list of names loads only those reaches. | +| `quantities` | `None` \| `str` \| `list[str]` | Control which quantities are read at each selected location. `None` (default) reads everything; `[]` reads nothing; a name or list reads only those quantities. | ::: {.callout-note} Selective loading only controls **which timeseries are held in memory**. The full network topology (nodes, reaches, lengths) is always constructed so that `find()`, `recall()`, and graph algorithms still work on the complete network. @@ -283,6 +284,13 @@ network_subset = Network.from_mike( network_subset ``` +`quantities` cuts the read down a third way, which helps when a calibration loop rebuilds the network for every trial but only scores one quantity. A location that does not carry a requested quantity simply stays topology-only: + +```{python} +network_discharge = Network.from_mike(path_to_res1d, quantities="Discharge") +network_discharge.quantities +``` + When only some nodes are loaded, `to_dataframe()` and `to_dataset()` only contain columns for those nodes — the rest are graph-connected but data-free: ```{python} diff --git a/src/modelskill/model/adapters/_res1d.py b/src/modelskill/model/adapters/_res1d.py index 2c335b317..3ed0032e5 100644 --- a/src/modelskill/model/adapters/_res1d.py +++ b/src/modelskill/model/adapters/_res1d.py @@ -16,7 +16,9 @@ _EMPTY_DATA = pd.DataFrame() -def _simplify_colnames(node: ResultNode | ResultGridPoint) -> pd.DataFrame: +def _simplify_colnames( + node: ResultNode | ResultGridPoint, quantities: set[str] | None = None +) -> pd.DataFrame: # We remove suffixes and indexes so the columns contain only the quantity names # Some formats keep no timeseries at all on some locations - MIKE 11, for instance, @@ -28,9 +30,24 @@ def _simplify_colnames(node: ResultNode | ResultGridPoint) -> pd.DataFrame: # The columns in a Res1D dataframe follow the convention "Quantity:Location:Sublocation" # where Location refers to the node id or the reach id followed by the chainage. RES1D_NAME_SEP = ":" - df = node.to_dataframe() + + available = list(node.quantities) + wanted = ( + available if quantities is None else [q for q in available if q in quantities] + ) + + if not wanted: + # A location need not carry every requested quantity; it stays topology-only. + return _EMPTY_DATA + + if len(wanted) == len(available): + # Reading the whole location is one interop call rather than one per quantity. + df = node.to_dataframe() + else: + df = pd.concat([getattr(node, q).to_dataframe() for q in wanted], axis=1) + renamer_dict = {} - for quantity in node.quantities: + for quantity in wanted: column_pairs = [ (col, quantity) for col in df.columns @@ -135,6 +152,7 @@ def __init__( *, populate_gridpoints: bool = True, length: float | None = None, + quantities: set[str] | None = None, ): self._id = reach.name @@ -169,7 +187,9 @@ def __init__( GridPoint( gridpoint.reach_name, gridpoint.chainage, - _simplify_colnames(gridpoint) if populate_gridpoints else None, + _simplify_colnames(gridpoint, quantities) + if populate_gridpoints + else None, ) for gridpoint in intermediate_gridpoints ] diff --git a/src/modelskill/network.py b/src/modelskill/network.py index 90bc3d848..ee0506081 100644 --- a/src/modelskill/network.py +++ b/src/modelskill/network.py @@ -422,6 +422,7 @@ def from_mike( *, nodes: str | list[str] | None = None, reaches: str | list[str] | None = None, + quantities: str | list[str] | None = None, ) -> Network: """Create a Network from a MIKE 1D or MIKE 11 result file. @@ -448,6 +449,17 @@ def from_mike( * A single reach name or a list of reach names — only those reaches get gridpoint data; others are topology-only. * ``[]`` (empty list) — no gridpoint data is loaded at all. + quantities : str, list of str, or None, optional + Controls which quantities are read at each selected location. + + * ``None`` *(default)* — every quantity is read. + * A single quantity name or a list of names — only those are read. + * ``[]`` (empty list) — no data is read at all. + + A location that does not carry a requested quantity becomes + topology-only rather than an error, so this composes with ``nodes`` + and ``reaches`` on files where nodes and reaches hold different + quantities. Returns ------- @@ -484,6 +496,11 @@ def from_mike( ... reaches=["reach_1"], ... ) + Read a single quantity, for a calibration loop that only scores + discharge: + + >>> network = Network.from_mike("model.res1d", quantities="Discharge") + Notes ----- MIKE 11 keeps its timeseries on reach gridpoints rather than on nodes, @@ -498,6 +515,7 @@ def from_mike( res, nodes=nodes, reaches=reaches, + quantities=quantities, allowed=_MIKE_EXTENSIONS, caller="from_mike", ) @@ -511,6 +529,7 @@ def from_epanet( inp: str | Path | None = None, nodes: str | list[str] | None = None, reaches: str | list[str] | None = None, + quantities: str | list[str] | None = None, ) -> Network: """Create a Network from an EPANET result file and its companions. @@ -539,6 +558,9 @@ def from_epanet( Which reaches get their gridpoint data loaded. See :meth:`from_mike`. EPANET results have no intermediate gridpoints, so this argument has no effect. + quantities : str, list of str, or None, optional + Which quantities are read at each selected location. See + :meth:`from_mike`. Returns ------- @@ -596,6 +618,7 @@ def from_epanet( res, nodes=nodes, reaches=reaches, + quantities=quantities, allowed=_EPANET_EXTENSIONS, caller="from_epanet", resx=resx, @@ -611,6 +634,7 @@ def _from_mikeio1d( reaches: str | list[str] | None, allowed: frozenset[str], caller: str, + quantities: str | list[str] | None = None, resx: str | Path | Res1D | None = None, inp: str | Path | None = None, ) -> Network: @@ -664,8 +688,22 @@ def _from_mikeio1d( extra = None if resx is None else cls._open_companion_result(res, resx) lengths = None if inp is None else cls._read_companion_lengths(inp) + # None is threaded through as "read everything" rather than expanded to + # res.quantities, which would only cost a lookup for the same result. + if quantities is None: + quantities_set: set[str] | None = None + elif isinstance(quantities, str): + quantities_set = {quantities} + else: + quantities_set = set(quantities) + list_of_reaches = cls._load_res1d_network( - res, nodes_list, reaches_list, extra=extra, lengths=lengths + res, + nodes_list, + reaches_list, + extra=extra, + lengths=lengths, + quantities=quantities_set, ) return cls(list_of_reaches) @@ -796,6 +834,7 @@ def _load_res1d_network( *, extra: Res1D | None = None, lengths: dict[str, float] | None = None, + quantities: set[str] | None = None, ) -> list[Res1DReach]: from modelskill.model.adapters._res1d import ( Res1DReach, @@ -822,16 +861,18 @@ def _init_node(reach: ResultReach, is_end: bool) -> Res1DNode: gpt_idx = -1 if is_end else 0 if id in nodes_set: if id not in node_data: - df = _simplify_colnames(res.nodes[id]) + df = _simplify_colnames(res.nodes[id], quantities) # Merged here rather than up front so selective loading still # decides what is held in memory. if extra is not None and id in extra.nodes: df = _merge_extra_quantities( - df, _simplify_colnames(extra.nodes[id]), node_id=id + df, + _simplify_colnames(extra.nodes[id], quantities), + node_id=id, ) node_data[id] = df overlapping_gridpoint = reach.gridpoints[gpt_idx] - boundary = _simplify_colnames(overlapping_gridpoint) + boundary = _simplify_colnames(overlapping_gridpoint, quantities) return Res1DNode( id, data=node_data[id], boundary={reach.name: boundary} ) @@ -845,6 +886,7 @@ def _init_node(reach: ResultReach, is_end: bool) -> Res1DNode: _init_node(reach, True), populate_gridpoints=reach.name in reaches_set, length=lengths.get(reach.name), + quantities=quantities, ) for reach in res.reaches.values() ] diff --git a/tests/test_network.py b/tests/test_network.py index 54d249d24..ce388e647 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -876,6 +876,104 @@ def test_res1d_opened_with_a_path_is_refused(self): Network.from_mike(res) +# --------------------------------------------------------------------------- +# from_mike — quantities filter +# +# In network.res1d every node carries WaterLevel only, interior gridpoints +# carry either Discharge or WaterLevel. +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + sys.version_info >= (3, 14), reason="mikeio1d requires Python < 3.14" +) +def test_from_mike_reads_all_quantities_by_default(): + path_to_file = "./tests/testdata/network.res1d" + network = Network.from_mike(path_to_file) + + assert set(network.quantities) == {"WaterLevel", "Discharge"} + + +@pytest.mark.skipif( + sys.version_info >= (3, 14), reason="mikeio1d requires Python < 3.14" +) +def test_from_mike_quantities_filter_reads_only_requested_quantity(): + path_to_file = "./tests/testdata/network.res1d" + network = Network.from_mike(path_to_file, quantities="Discharge") + + assert network.quantities == ["Discharge"] + + +@pytest.mark.skipif( + sys.version_info >= (3, 14), reason="mikeio1d requires Python < 3.14" +) +def test_from_mike_quantities_filter_leaves_other_locations_topology_only(): + """A location that does not carry the requested quantity is not an error.""" + path_to_file = "./tests/testdata/network.res1d" + network = Network.from_mike(path_to_file, quantities="Discharge") + g = network.graph + + assert g.nodes[network.find(node="1")]["data"].empty + assert ( + g.number_of_nodes() == Network.from_mike(path_to_file).graph.number_of_nodes() + ) + + +@pytest.mark.skipif( + sys.version_info >= (3, 14), reason="mikeio1d requires Python < 3.14" +) +def test_from_mike_quantities_filter_populates_matching_breakpoints(): + path_to_file = "./tests/testdata/network.res1d" + network = Network.from_mike(path_to_file, quantities="Discharge") + + breakpoints = network._reaches["100l1"].breakpoints + populated = [bp for bp in breakpoints if not bp.data.empty] + + assert [bp.quantities for bp in populated] == [["Discharge"]] + + +@pytest.mark.skipif( + sys.version_info >= (3, 14), reason="mikeio1d requires Python < 3.14" +) +def test_from_mike_quantities_filter_keeps_node_data_for_node_quantity(): + path_to_file = "./tests/testdata/network.res1d" + network = Network.from_mike(path_to_file, quantities="WaterLevel") + + assert network.quantities == ["WaterLevel"] + assert not network.graph.nodes[network.find(node="1")]["data"].empty + + +@pytest.mark.skipif( + sys.version_info >= (3, 14), reason="mikeio1d requires Python < 3.14" +) +def test_from_mike_empty_quantities_keeps_topology_and_loads_nothing(): + path_to_file = "./tests/testdata/network.res1d" + full_network = Network.from_mike(path_to_file) + network = Network.from_mike(path_to_file, quantities=[]) + + assert network.graph.number_of_nodes() == full_network.graph.number_of_nodes() + assert network.quantities == [] + + df = network.to_dataframe() + assert df.empty + assert isinstance(df.columns, pd.MultiIndex) + + +@pytest.mark.skipif( + sys.version_info >= (3, 14), reason="mikeio1d requires Python < 3.14" +) +def test_from_mike_quantities_filter_combines_with_nodes_filter(): + path_to_file = "./tests/testdata/network.res1d" + network = Network.from_mike( + path_to_file, nodes=["1"], reaches=[], quantities=["WaterLevel"] + ) + g = network.graph + + nodes_with_data = [n for n in g.nodes if not g.nodes[n]["data"].empty] + assert [network.recall(n)["node"] for n in nodes_with_data] == ["1"] + assert list(g.nodes[network.find(node="1")]["data"].columns) == ["WaterLevel"] + + # --------------------------------------------------------------------------- # NodeObservation — alias / breakpoint node forms # --------------------------------------------------------------------------- From 4c39ce0e4f8a4a6dd3631cf067e82ef2a57d8b59 Mon Sep 17 00:00:00 2001 From: jpalm3r Date: Tue, 11 Aug 2026 10:19:30 +0200 Subject: [PATCH 3/3] Add drivers for measuring network load times Two scripts behind the work above: one profiles a single load, the other compares load times across revisions so a claimed speed-up can be checked rather than asserted. Add snakeviz for reading the profiles, and ignore the generated output directory. Co-Authored-By: Claude Opus 5 --- .gitignore | 4 +- pyproject.toml | 9 +- tests/profiling/compare_network_loading.py | 277 +++++++++++++++++++++ tests/profiling/profile_network_loading.py | 207 +++++++++++++++ 4 files changed, 495 insertions(+), 2 deletions(-) create mode 100644 tests/profiling/compare_network_loading.py create mode 100644 tests/profiling/profile_network_loading.py diff --git a/.gitignore b/.gitignore index b38653822..6b4b19d5c 100644 --- a/.gitignore +++ b/.gitignore @@ -155,4 +155,6 @@ docs/api/*.qmd uv.lock -tests/testdata/confidential/* \ No newline at end of file +tests/testdata/confidential/* + +tests/profiling/output/ \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 890fc7574..925c4485b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,7 +46,14 @@ classifiers = [ networks = ["mikeio1d", "networkx"] [dependency-groups] -dev = ["pytest", "plotly >= 4.5", "ruff==0.6.2", "netCDF4", "dask"] +dev = [ + "pytest", + "plotly >= 4.5", + "ruff==0.6.2", + "netCDF4", + "dask", + "snakeviz>=2.2.2", +] docs = ["quartodoc==0.11.1", "nbformat", "nbconvert", "ipykernel", "griffe<2"] diff --git a/tests/profiling/compare_network_loading.py b/tests/profiling/compare_network_loading.py new file mode 100644 index 000000000..3305c2360 --- /dev/null +++ b/tests/profiling/compare_network_loading.py @@ -0,0 +1,277 @@ +"""Time `Network.from_epanet` across revisions, file sets and quantity filters. + +Each case runs as its own subprocess of `profile_network_loading.py`, so no +module-level cache carries between them. A case can be pointed at a different +modelskill source tree with `--baseline-src`, which is how this compares the +current checkout against another revision: + + git worktree add ../modelskill-main main + + uv run python tests/profiling/compare_network_loading.py \\ + --res-path model.res --resx model.resx --inp model.inp \\ + --baseline-src ../modelskill-main/src --quantities Pressure + +Both revisions then run in the one virtual environment, with only the +modelskill source differing - mikeio1d and networkx stay pinned to the same +builds, so the numbers reflect modelskill's own code and nothing else. + +The `--baseline-src` tree need not know about the `quantities` filter: the +filtered cases only ever run against the current checkout. +""" + +import argparse +import json +import os +import statistics +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path + +RUNNER = Path(__file__).resolve().parent / "profile_network_loading.py" + + +@dataclass(frozen=True) +class Case: + """One row of the comparison.""" + + variant: str + files: str + src: Path | None + companions: bool + quantities: list[str] | None + + @property + def label(self) -> str: + return f"{self.variant}--{self.files}".replace(" ", "").replace("+", "-") + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--res-path", type=Path, required=True) + parser.add_argument( + "--resx", + type=Path, + default=None, + help="Companion '.resx'. Given together with --inp, a second file set " + "is timed alongside the res-only one.", + ) + parser.add_argument("--inp", type=Path, default=None) + parser.add_argument( + "--baseline-src", + type=Path, + default=None, + help="src/ directory of another modelskill checkout - typically a " + "worktree of main - to time as the baseline. Omit to time only the " + "current checkout.", + ) + parser.add_argument( + "--baseline-name", + default="main", + help="What to call the --baseline-src revision in the table.", + ) + parser.add_argument( + "--current-name", + default="current", + help="What to call the current checkout in the table, e.g. its branch name.", + ) + parser.add_argument( + "--quantities", + nargs="+", + default=None, + help="Quantity name(s) for the filtered variant, e.g. Pressure. Omit " + "to skip the filtered cases.", + ) + parser.add_argument( + "--rounds", + type=int, + default=3, + help="How many times to run the whole case list. Each measurement is " + "one load in a fresh process - loading this much data repeatedly in " + "one process builds memory pressure that swamps the differences being " + "measured. The reported figure is the minimum over the rounds.", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Also dump a cProfile pass per case, for drilling into a " + "surprising number afterwards. Roughly doubles the runtime.", + ) + return parser.parse_args() + + +def _build_cases(args: argparse.Namespace) -> list[Case]: + companions = [n for n, p in (("resx", args.resx), ("inp", args.inp)) if p] + file_sets = [("res", False)] + if companions: + file_sets.append(("res+" + "+".join(companions), True)) + + variants: list[tuple[str, Path | None, list[str] | None]] = [] + if args.baseline_src is not None: + variants.append((args.baseline_name, args.baseline_src.resolve(), None)) + variants.append((args.current_name, None, None)) + if args.quantities: + filtered = f"{args.current_name} +{','.join(args.quantities)}" + variants.append((filtered, None, args.quantities)) + + return [ + Case(variant, files, src, companions, quantities) + for variant, src, quantities in variants + for files, companions in file_sets + ] + + +def _run_once(case: Case, args: argparse.Namespace, *, profile: bool) -> dict | None: + cmd = [ + sys.executable, + str(RUNNER), + "--res-path", + str(args.res_path), + # One load, no in-process warm-up. Every measurement therefore includes + # the same one-off mikeio1d start-up cost, which is what a user pays + # too, and no measurement is distorted by an earlier load's memory. + "--repeat", + "1", + "--warmup", + "0", + "--label", + case.label, + "--json", + ] + if case.companions: + if args.resx is not None: + cmd += ["--resx", str(args.resx)] + if args.inp is not None: + cmd += ["--inp", str(args.inp)] + if case.quantities: + cmd += ["--quantities", *case.quantities] + if profile: + cmd += ["--profile"] + + env = dict(os.environ) + if case.src is None: + env.pop("PYTHONPATH", None) + else: + # Wins over the .pth entry that site-packages adds for the editable + # install, so this subprocess imports the baseline tree instead. + env["PYTHONPATH"] = str(case.src) + + print(f" {case.variant:<28} {case.files:<14} ... ", end="", flush=True) + proc = subprocess.run(cmd, env=env, capture_output=True, text=True) + if proc.returncode != 0: + print("FAILED") + print(proc.stdout, proc.stderr, sep="\n", file=sys.stderr) + return None + + result = json.loads(proc.stdout.strip().splitlines()[-1]) + print(f"{result['min']:.2f}s") + return {**result, "variant": case.variant, "files": case.files} + + +def _table(results: list[dict], baseline_name: str) -> str: + """CPU time leads, because it is the column the machine cannot distort.""" + by_files = {r["files"]: r for r in results if r["variant"] == baseline_name} + + lines = [ + "| revision | files | cpu min | cpu speedup | wall min | wall speedup |", + "| --- | --- | ---: | ---: | ---: | ---: |", + ] + for r in results: + ref = by_files.get(r["files"]) + cpu_up = "-" if ref is None else f"{ref['cpu_min'] / r['cpu_min']:.2f}x" + wall_up = "-" if ref is None else f"{ref['min'] / r['min']:.2f}x" + lines.append( + f"| {r['variant']} | {r['files']} | {r['cpu_min']:.2f}s | {cpu_up} | " + f"{r['min']:.2f}s | {wall_up} |" + ) + return "\n".join(lines) + + +def _rounds(results: list[dict]) -> str: + lines = [ + "| revision | files | cpu per round | wall per round |", + "| --- | --- | --- | --- |", + ] + for r in results: + cpu = ", ".join(f"{s:.1f}" for s in r["cpu_samples"]) + wall = ", ".join(f"{s:.1f}" for s in r["samples"]) + lines.append(f"| {r['variant']} | {r['files']} | {cpu} | {wall} |") + return "\n".join(lines) + + +def _sanity(results: list[dict]) -> str: + """Different node or reach counts would mean the cases are not comparable.""" + lines = [ + "| revision | files | nodes | reaches | quantities loaded | source |", + "| --- | --- | ---: | ---: | --- | --- |", + ] + for r in results: + lines.append( + f"| {r['variant']} | {r['files']} | {r['n_nodes']} | {r['n_reaches']} | " + f"{', '.join(r['quantities'])} | {r['modelskill']} |" + ) + return "\n".join(lines) + + +def main() -> None: + args = _parse_args() + cases = _build_cases(args) + + print(f"Timing {len(cases)} cases over {args.rounds} rounds.") + + samples: dict[str, list[float]] = {c.label: [] for c in cases} + cpu_samples: dict[str, list[float]] = {c.label: [] for c in cases} + latest: dict[str, dict] = {} + for round_ in range(args.rounds): + print(f"\nRound {round_ + 1}/{args.rounds}") + # Reversed on alternate rounds so a machine that drifts over the run + # biases the first and last case equally rather than only the last. + ordered = cases if round_ % 2 == 0 else list(reversed(cases)) + for case in ordered: + profile = args.profile and round_ == args.rounds - 1 + result = _run_once(case, args, profile=profile) + if result is None: + continue + samples[case.label].append(result["seconds"][0]) + cpu_samples[case.label].append(result["cpu_seconds"][0]) + latest[case.label] = result + + results = [ + { + **latest[c.label], + "samples": samples[c.label], + "cpu_samples": cpu_samples[c.label], + "min": min(samples[c.label]), + "median": statistics.median(samples[c.label]), + "cpu_min": min(cpu_samples[c.label]), + } + for c in cases + if samples[c.label] + ] + if not results: + sys.exit("Every case failed.") + + print("\n" + _table(results, args.baseline_name)) + print("\n" + _rounds(results)) + print("\n" + _sanity(results)) + + counts = {(r["n_nodes"], r["n_reaches"]) for r in results} + if len(counts) > 1: + print( + "\nWARNING: the cases loaded different networks " + f"({counts}), so the times are not comparable." + ) + + # A wide spread within a case means the machine was noisy and the + # between-case differences above deserve less trust. + for name, key in (("cpu", "cpu_samples"), ("wall", "samples")): + spread = [max(r[key]) / min(r[key]) for r in results] + print( + f"\nWithin-case {name} spread (max/min): " + f"{statistics.mean(spread):.2f}x mean, {max(spread):.2f}x worst. " + "1.00x is perfectly repeatable." + ) + + +if __name__ == "__main__": + main() diff --git a/tests/profiling/profile_network_loading.py b/tests/profiling/profile_network_loading.py new file mode 100644 index 000000000..38836efc6 --- /dev/null +++ b/tests/profiling/profile_network_loading.py @@ -0,0 +1,207 @@ +"""Time one `Network.from_epanet` configuration, optionally under cProfile. + +This runs a single case. To compare several at once - different branches, +different companion files, filtered against unfiltered - use the sibling +`compare_network_loading.py`, which drives this script. + +Usage +----- +uv run python tests/profiling/profile_network_loading.py --res-path + +Companion files and a quantities filter are opt-in: +uv run python tests/profiling/profile_network_loading.py --res-path --resx --inp +uv run python tests/profiling/profile_network_loading.py --res-path --quantities Pressure + +Add --profile for a cProfile pass on top of the timed runs, then inspect it with: +uv run snakeviz tests/profiling/output/profile.prof + +An argument is only forwarded to `from_epanet` when it is given, so this script +also runs against revisions that predate one of its parameters. That is what +lets the same file time `main` and a feature branch. +""" + +import argparse +import cProfile +import gc +import json +import pstats +import statistics +import sys +import time +from pathlib import Path + +import modelskill +from modelskill.network import Network + +DEFAULT_OUTPUT_DIR = Path(__file__).resolve().parent / "output" + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--res-path", + type=Path, + required=True, + help="Path to the EPANET .res file to load.", + ) + parser.add_argument( + "--resx", + type=Path, + default=None, + help="Optional companion '.resx' file, merged in via from_epanet's " + "resx= parameter.", + ) + parser.add_argument( + "--inp", + type=Path, + default=None, + help="Optional companion '.inp' file, read for reach lengths via " + "from_epanet's inp= parameter.", + ) + parser.add_argument( + "--quantities", + nargs="+", + default=None, + help="Quantity name(s) to pass as the 'quantities' filter to " + "Network.from_epanet, e.g. --quantities Pressure. Omit to load " + "every quantity (the default, unfiltered behavior).", + ) + parser.add_argument( + "--repeat", + type=int, + default=3, + help="How many timed loads to run. The reported figures are the " + "minimum and median over these.", + ) + parser.add_argument( + "--warmup", + type=int, + default=1, + help="Loads to run before timing starts. The first load in a process " + "pays a one-off mikeio1d/.NET start-up cost of several seconds, which " + "is identical across revisions and would otherwise swamp the " + "comparison. Its duration is still reported, as 'cold'.", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Run one extra load under cProfile and dump it to " + "/