diff --git a/LoopStructural/modelling/core/_model_exporter.py b/LoopStructural/modelling/core/_model_exporter.py index 8442f6b51..c63305642 100644 --- a/LoopStructural/modelling/core/_model_exporter.py +++ b/LoopStructural/modelling/core/_model_exporter.py @@ -46,7 +46,7 @@ def get_stratigraphic_surfaces( units = [] if model.stratigraphic_column is None: return [] - units = model.stratigraphic_column.get_isovalues() + units = model.stratigraphic_column.get_isovalues(where='bottom' if bottoms else 'top') units_for_group = {} for name, u in units.items(): if u['group'] not in model: diff --git a/LoopStructural/modelling/core/stratigraphic_column.py b/LoopStructural/modelling/core/stratigraphic_column.py index 943c68ba4..a2d1971f1 100644 --- a/LoopStructural/modelling/core/stratigraphic_column.py +++ b/LoopStructural/modelling/core/stratigraphic_column.py @@ -677,16 +677,29 @@ def from_dict(cls, data): column.add_element(element) return column - def get_isovalues(self) -> dict[str, float]: + def get_isovalues(self, where: str = 'bottom') -> dict[str, float]: """ Returns a dictionary of isovalues for the stratigraphic units in the column. + + Parameters + ---------- + where : str, optional + 'bottom' (default) returns the value at the base of each unit. + 'top' returns the value at the top of each unit. """ + if where not in ('top', 'bottom'): + raise ValueError("Invalid 'where' argument. Use 'top' or 'bottom'.") surface_values = {} for g in reversed(self.get_groups()): v = 0 for u in reversed(g.units): - surface_values[u.name] = {'value': v, 'group': g.name, 'colour': u.colour} + base = v v += u.thickness + surface_values[u.name] = { + 'value': v if where == 'top' else base, + 'group': g.name, + 'colour': u.colour, + } return surface_values def plot(self, *, ax=None, **kwargs): diff --git a/tests/unit/modelling/test_stratigraphic_column.py b/tests/unit/modelling/test_stratigraphic_column.py index 1f7fc2849..db27518b8 100644 --- a/tests/unit/modelling/test_stratigraphic_column.py +++ b/tests/unit/modelling/test_stratigraphic_column.py @@ -444,6 +444,25 @@ def test_get_isovalues_multi_unit_group(self): assert isovalues["B"]["value"] == 10 assert isovalues["C"]["value"] == 15 + def test_get_isovalues_where_top_returns_unit_top(self): + # where='top' must give each unit's top, i.e. base + thickness -- + # the same value that where='bottom' (default) gives the next + # (younger) unit as its base. + column = StratigraphicColumn() + column.clear(basement=False) + column.add_unit("A", thickness=10, id=0) + column.add_unit("B", thickness=5, id=1) + column.add_unit("C", thickness=3, id=2) + isovalues = column.get_isovalues(where='top') + assert isovalues["A"]["value"] == 10 + assert isovalues["B"]["value"] == 15 + assert isovalues["C"]["value"] == 18 + + def test_get_isovalues_invalid_where_raises(self): + column = self._build_two_group_column() + with pytest.raises(ValueError): + column.get_isovalues(where="middle") + class TestOrderingAndUpdates: def test_update_order_reorders_elements(self):