Merge pull request #29 from GNS-Science/fix/filtered-branches

Fix/filtered branches

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.5.0
+current_version = 0.5.1
 commit = True
 tag = True
 

CHANGELOG.md

@@ -1,7 +1,12 @@
 # Changelog
 
+## [0.5.1] 2023-12-14
+### Changed
+ - SourceLogicTree.from_branches() returns logic tree with Branch objects rather than FilterdBranch objects
+ - remove whitespace from logic tree file paths for compatability with OpenQuake
+
 ## [0.5.0] 2023-12-12
-### Added
+## Added
  - support caching of downloads
  - build sources xml
  - migration of version1 to version2 SLT

nzshm_model/__init__.py

@@ -1,7 +1,7 @@
 from . import nshm_v1_0_0, nshm_v1_0_4
 
 # Python package version is different than the NSHM MODEL version !!
-__version__ = '0.5.0'
+__version__ = '0.5.1'
 
 CURRENT_VERSION = "NSHM_v1.0.0"
 

nzshm_model/psha_adapter/openquake/logic_tree.py

@@ -38,6 +38,7 @@
     GroundMotionUncertaintyModel,
     NshmSourceUncertaintyModel,
     SourcesUncertaintyModel,
+    _strip_whitespace,
 )
 
 NRML_NS = None
@@ -94,7 +95,7 @@ def from_parent_slt(
             yield _instance
 
     def path(self) -> PurePath:
-        return PurePath(self.parent.path(), self.branchID)
+        return PurePath(self.parent.path(), _strip_whitespace(self.branchID))
 
 
 @dataclass
@@ -140,7 +141,7 @@ def uncertainty_class(self):
         return GenericUncertaintyModel
 
     def path(self) -> PurePath:
-        return PurePath(self.parent.path(), self.branchSetID)
+        return PurePath(self.parent.path(), _strip_whitespace(self.branchSetID))
 
 
 @dataclass
@@ -156,7 +157,7 @@ def from_parent_element(cls, root: objectify.Element) -> Iterator["LogicTree"]:
             yield _instance
 
     def path(self) -> PurePath:
-        return PurePath(self.logicTreeID)
+        return PurePath(_strip_whitespace(self.logicTreeID))
 
     @classmethod
     def from_parent_slt(cls, slt: "slt.SourceLogicTree") -> "LogicTree":

nzshm_model/psha_adapter/openquake/simple_nrml.py

@@ -1,7 +1,7 @@
 import logging
 import pathlib
 import zipfile
-from typing import Dict, Union
+from typing import Any, Dict, Generator, Union
 
 from lxml import etree
 from lxml.builder import ElementMaker
@@ -95,7 +95,7 @@ def build_sources_xml(self, source_map):
                         for filepath in source_map.get(source.nrml_id):
                             if not filepath.suffix == '.xml':
                                 continue
-                            files += f"\t'{filepath}'\n"
+                            files += f"\t{filepath}\n"
                         ltb = LTB(UM(files), UW(str(branch.weight)), branchID=branch_name)
                 # else:
                 #     # old style logic tree
@@ -118,7 +118,7 @@ def write_config(
         cache_folder: Union[pathlib.Path, str],
         target_folder: Union[pathlib.Path, str],
         source_map: Union[None, Dict[str, list[pathlib.Path]]] = None,
-    ):  # -> pathlib.Path
+    ) -> pathlib.Path:
         destination = pathlib.Path(target_folder)
         assert destination.exists()
         assert destination.is_dir()
@@ -129,9 +129,11 @@ def write_config(
         target_file = pathlib.Path(destination, 'sources.xml')
         with open(target_file, 'w') as fout:
             fout.write(xmlstr)
-        return True
+        return target_file
 
-    def unpack_resources(self, cache_folder: Union[pathlib.Path, str], target_folder: Union[pathlib.Path, str]):
+    def unpack_resources(
+        self, cache_folder: Union[pathlib.Path, str], target_folder: Union[pathlib.Path, str]
+    ) -> Dict[str, list[pathlib.Path]]:
         target = pathlib.Path(target_folder)
         target.mkdir(parents=True, exist_ok=True)
         source_map = {}
@@ -163,15 +165,17 @@ def unpack_resources(self, cache_folder: Union[pathlib.Path, str], target_folder
         #         prefixed = pathlib.Path(destination, f"{file_prefix}_{name}")
         #         extracted.rename(prefixed)
 
-    def fetch_resources(self, cache_folder: Union[pathlib.Path, str]):
+    def fetch_resources(
+        self, cache_folder: Union[pathlib.Path, str]
+    ) -> Generator[tuple[Any, pathlib.Path, Any], None, None]:
         destination = pathlib.Path(cache_folder)
         destination.mkdir(parents=True, exist_ok=True)
         nrml_logic_tree = self.config()
         for branch_set in nrml_logic_tree.branch_sets:
             for branch in branch_set.branches:
                 for um in branch.uncertainty_models:
                     filepath = fetch_toshi_source(um.toshi_nrml_id, destination)
-                    yield tuple([um.toshi_nrml_id, filepath, um])
+                    yield um.toshi_nrml_id, filepath, um
 
     def config(self):
         return NrmlDocument.from_model_slt(self._source_logic_tree).logic_trees[0]

nzshm_model/psha_adapter/openquake/uncertainty_models.py

@@ -13,6 +13,10 @@
     from .logic_tree import LogicTreeBranch
 
 
+def _strip_whitespace(string: str) -> str:
+    return ''.join(string.split())
+
+
 @dataclass
 class GenericUncertaintyModel:
     parent: "LogicTreeBranch"
@@ -23,7 +27,7 @@ def from_parent_element(cls, node: objectify.Element, parent: "LogicTreeBranch")
         return GenericUncertaintyModel(parent=parent, text=node.text.strip())
 
     def path(self) -> PurePath:
-        return PurePath(self.parent.path(), self.text)
+        return PurePath(self.parent.path(), _strip_whitespace(self.text))
 
 
 @dataclass
@@ -44,7 +48,7 @@ def from_parent_element(cls, node: objectify.Element, parent: "LogicTreeBranch")
         )
 
     def path(self) -> PurePath:
-        return PurePath(self.parent.path(), self.text)  # todo unique combination of name + args
+        return PurePath(self.parent.path(), _strip_whitespace(self.text))  # todo unique combination of name + args
 
 
 @dataclass
@@ -61,7 +65,7 @@ def from_parent_element(cls, node: objectify.Element, parent: "LogicTreeBranch")
         )
 
     def path(self) -> PurePath:
-        return PurePath(self.parent.path(), self.text)  # todo unique combination of sources
+        return PurePath(self.parent.path(), _strip_whitespace(self.text))  # todo unique combination of sources
 
 
 @dataclass
@@ -87,4 +91,4 @@ def from_parent_slt(cls, ltb: "slt.Branch", parent: "LogicTreeBranch") -> Iterat
                 )
 
     def path(self) -> PurePath:
-        return PurePath(self.parent.path(), self.toshi_nrml_id)  # todo unique combination of sources
+        return PurePath(self.parent.path(), _strip_whitespace(self.toshi_nrml_id))  # todo unique combination of sources

nzshm_model/source_logic_tree/version2/logic_tree.py

@@ -7,7 +7,7 @@
 import json
 import pathlib
 from dataclasses import dataclass, field
-from typing import Dict, List, Type, Union
+from typing import Dict, Iterator, List, Type, Union
 
 import dacite
 
@@ -46,12 +46,6 @@ def tag(self):
         return str(self.values)
 
 
-@dataclass
-class FilteredBranch(Branch):
-    fslt: Union['FaultSystemLogicTree', None] = None  # this should never be serialised, only used for filtering
-    slt: Union['SourceLogicTree', None] = None  # this should never be serialised, only used for filtering
-
-
 @dataclass
 class FaultSystemLogicTree(FaultSystemLogicTreeBase):
     branches: List[Branch] = field(default_factory=list)
@@ -156,7 +150,7 @@ def __next__(self):
             return self.__branch_list[self.__current_branch - 1]
 
     @staticmethod
-    def from_branches(branches):
+    def from_branches(branches: Iterator['FilteredBranch']) -> 'SourceLogicTree':
         """
         Build a complete SLT from a iterable od branches.
 
@@ -168,18 +162,30 @@ def match_fslt(slt: SourceLogicTree, fb):
                 if fb.fslt.short_name == fslt.short_name:
                     return fslt
 
-        slt = None
+        version = None
         for fb in branches:
             # ensure an slt
-            if not slt:
+            if not version:
                 slt = SourceLogicTree(version=fb.slt.version, title=fb.slt.title)
+                version = fb.slt.version
             else:
-                assert slt.version == fb.slt.version
+                assert version == fb.slt.version
 
             # ensure an fslt
             fslt = match_fslt(slt, fb)
             if not fslt:
                 fslt = FaultSystemLogicTree(short_name=fb.fslt.short_name, long_name=fb.fslt.long_name)
                 slt.fault_systems.append(fslt)
-            fslt.branches.append(fb)
+            fslt.branches.append(fb.to_branch())
         return slt
+
+
+# this should never be serialised, only used for filtering
+@dataclass
+class FilteredBranch(Branch):
+    fslt: 'FaultSystemLogicTree' = FaultSystemLogicTree('shortname', 'longname')
+    slt: 'SourceLogicTree' = SourceLogicTree('version', 'title')
+
+    def to_branch(self) -> Branch:
+        branch_attributes = {k: v for k, v in self.__dict__.items() if k not in ('fslt', 'slt')}
+        return Branch(**branch_attributes)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nzshm-model"
-version = "0.5.0"
+version = "0.5.1"
 description = "The logic tree definitions, final configurations, and versioning of the New Zealand | Aotearoa National Seismic Hazard Model"
 authors = ["Chris DiCaprio <[email protected]>", "Chris Chamberlain <[email protected]>"]
 license = "AGPL3"

tests/psha_adapter/openquake/test_logic_tree_from_model.py

@@ -52,7 +52,7 @@ def test_source_logic_tree():
     assert src_logic_tree.branch_sets[0].branches[0].uncertainty_weight == pytest.approx(0.210000)
     assert (
         src_logic_tree.branch_sets[0].branches[0].path()
-        == PurePath('SLT_v9p0p0') / "PUY" / "[dm0.7, bN[0.902, 4.6], C4.0, s0.28]"
+        == PurePath('SLT_v9p0p0') / "PUY" / "[dm0.7,bN[0.902,4.6],C4.0,s0.28]"
     )
 
 
@@ -67,7 +67,7 @@ def test_source_logic_tree_uncertainty_PUY():
 
     assert src_logic_tree.branch_sets[BSID].branchSetID == "PUY"
     assert src_logic_tree.branch_sets[BSID].branches[0].path() == PurePath(
-        'SLT_v9p0p0', 'PUY', '[dm0.7, bN[0.902, 4.6], C4.0, s0.28]'
+        'SLT_v9p0p0', 'PUY', '[dm0.7,bN[0.902,4.6],C4.0,s0.28]'
     )
 
     assert len(src_logic_tree.branch_sets[BSID].branches[0].uncertainty_models) == 2
@@ -101,7 +101,7 @@ def test_source_logic_tree_uncertainty_SLAB():
     BSID = 3
 
     assert src_logic_tree.branch_sets[BSID].branchSetID == "SLAB"
-    assert src_logic_tree.branch_sets[BSID].branches[0].path() == PurePath('SLT_v9p0p0', 'SLAB', '[runiform, d1]')
+    assert src_logic_tree.branch_sets[BSID].branches[0].path() == PurePath('SLT_v9p0p0', 'SLAB', '[runiform,d1]')
 
     assert len(src_logic_tree.branch_sets[BSID].branches[0].uncertainty_models) == 1
 
@@ -119,7 +119,7 @@ def test_source_logic_tree_uncertainty_CRU():
     BSID = 2
     assert src_logic_tree.branch_sets[BSID].branchSetID == "CRU"
     assert src_logic_tree.branch_sets[BSID].branches[0].path() == PurePath(
-        'SLT_v9p0p0/CRU/[dmgeodetic, tdFalse, bN[0.823, 2.7], C4.2, s0.66]'
+        'SLT_v9p0p0/CRU/[dmgeodetic,tdFalse,bN[0.823,2.7],C4.2,s0.66]'
     )
     assert len(src_logic_tree.branch_sets[BSID].branches[0].uncertainty_models) == 2
 

tests/psha_adapter/openquake/test_logic_tree_from_xml.py

@@ -5,6 +5,7 @@
 import pytest
 
 from nzshm_model.psha_adapter import NrmlDocument
+from nzshm_model.psha_adapter.openquake.uncertainty_models import _strip_whitespace
 
 FIXTURE_PATH = Path(__file__).parent / "fixtures"
 
@@ -56,15 +57,14 @@ def test_nrml_gmm_logic_tree_paths():
         "lt1",
         "bs_crust",
         "STF22_upper",
-        '[Stafford2022]\n                  mu_branch = "Upper"',
+        '[Stafford2022]mu_branch="Upper"',
     )
 
     assert doc.logic_trees[0].branch_sets[1].branches[0].uncertainty_models[0].path() == PurePath(
         "lt1",
         "bs_slab",
         "Kuehn2020SS_GLO_lower",
-        '[KuehnEtAl2020SSlab]\n                        region = "GLO"'
-        '\n                        sigma_mu_epsilon = -1.28155',
+        '[KuehnEtAl2020SSlab]region="GLO"' 'sigma_mu_epsilon=-1.28155',
     )
 
 
@@ -161,5 +161,8 @@ def test_nrml_srm_logic_tree_path(
     assert doc.logic_trees[0].branch_sets[0].branches[0].path() == PurePath(logic_tree_id, branch_set_id, branch_id)
 
     assert doc.logic_trees[0].branch_sets[0].branches[0].uncertainty_models[0].path() == PurePath(
-        logic_tree_id, branch_set_id, branch_id, uncertainty_model
+        _strip_whitespace(logic_tree_id),
+        _strip_whitespace(branch_set_id),
+        _strip_whitespace(branch_id),
+        _strip_whitespace(uncertainty_model),
     )

tests/test_slt_filtering.py

@@ -5,6 +5,7 @@
 
 import nzshm_model
 from nzshm_model.source_logic_tree import SourceLogicTree
+from nzshm_model.source_logic_tree.version2 import Branch
 
 
 ## three example filter functions
@@ -67,3 +68,4 @@ def test_build_slt_from_filtered_slt(full_slt):
     print(slt)
     assert len(slt.fault_systems) == 1
     assert len(slt.fault_systems[0].branches) == 9
+    assert type(slt.fault_systems[0].branches[0]) is Branch  # isinstance() not used to avoid true for inherited classes