From f2d632c2af9db7bcd8b382e586115bea96cb8708 Mon Sep 17 00:00:00 2001
From: hvasbath <Hannes.Bathke@gmx.net>
Date: Wed, 5 Jun 2024 23:52:05 +0200
Subject: [PATCH] bem.base: add get_geometry to BEMResonse for plotting in
 sparrow

---
 beat/apps/beat.py | 30 ++++++++++++++++++++++--------
 beat/bem/base.py  | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 8 deletions(-)

diff --git a/beat/apps/beat.py b/beat/apps/beat.py
index 38e6dd0f..fae450b3 100644
--- a/beat/apps/beat.py
+++ b/beat/apps/beat.py
@@ -2356,15 +2356,16 @@ def setup(parser):
     dump(rpoint, filename=outpoint_name)
     logger.info("Dumped %s solution to %s" % (options.post_llk, outpoint_name))
 
-    if options.mode == ffi_mode_str:
-        if "seismic" in problem.config.problem_config.datatypes:
-            datatype = "seismic"
-        elif "geodetic" in problem.config.problem_config.datatypes:
-            datatype = "geodetic"
-        else:
-            logger.info("Rupture geometry only available for static / kinematic data.")
+    # get geometry
+    if "seismic" in problem.config.problem_config.datatypes:
+        datatype = "seismic"
+    elif "geodetic" in problem.config.problem_config.datatypes:
+        datatype = "geodetic"
+    else:
+        logger.info("Rupture geometry only available for static / kinematic data.")
 
-        comp = problem.composites[datatype]
+    comp = problem.composites[datatype]
+    if options.mode == ffi_mode_str:
         engine = LocalEngine(store_superdirs=[comp.config.gf_config.store_superdir])
 
         target = comp.targets[0]
@@ -2395,6 +2396,19 @@ def setup(parser):
             )
             dump(geom, filename=ffi_rupture_table_path)
 
+    elif options.mode == bem_mode_str:
+
+        comp.point2sources(point)
+        response = composite.engine.process(
+            sources=composite.sources, targets=composite.targets
+        )
+
+        geoms = response.get_geometry(problem.event)
+        for k, geom in enumerate(geoms):
+            bem_geometry_path = pjoin(
+                results_path, "rupture_evolution_{}_{}.yaml".format(options.post_llk, k))
+            dump(geom, filename=bem_geometry_path)
+
     for datatype, composite in problem.composites.items():
         logger.info(
             "Exporting %s synthetics \n"
diff --git a/beat/bem/base.py b/beat/bem/base.py
index 15b2fe6b..1435fe61 100644
--- a/beat/bem/base.py
+++ b/beat/bem/base.py
@@ -98,6 +98,44 @@ def get_source_magnitudes(self, shear_modulus):
 
         return magnitudes
 
+    def get_geometry(self, event):
+        """Returns list of geometries of sources for plotting in the pyrocko.sparrow"""
+        def duplicate_property(array):
+            ndims = len(array.shape)
+            if ndims == 1:
+                return num.hstack((array, array))
+            elif ndims == 2:
+                return num.vstack((array, array))
+            else:
+                raise TypeError("Only 1-2d data supported!")
+
+        from pyrocko.model import Geometry
+        ncorners = 3
+
+        geoms = []
+        sources_slips = self.source_slips()
+        times = num.zeros(1)
+        for dsource, source_slips in zip(self.discretized_sources, sources_slips):
+            n_nodes = ncorners * dsource.n_vertices
+            latlon = num.ones((n_nodes, 2)) * num.array([event.lat, event.lon])
+            vertices = num.hstack([latlon, dsource.triangles_xyz.reshape((n_nodes, 3))])
+            geom = Geometry(times=times, event=event)
+
+            faces1 = num.arange(n_nodes, dtype="int64").reshape(
+                dsource.n_vertices, ncorners
+            )
+            faces2 = num.fliplr(faces1)
+            faces = num.vstack((faces1, faces2))
+
+            geom.setup(vertices, faces)
+            for slip_comp in ["strike", "dip", "normal"]:
+                slips = source_slips[slip_comp_to_idx[slip_comp]]
+                geom.add_property(((slip_comp, "float64", [])), duplicate_property(slips))
+
+            geoms.append(geom)
+
+        return geoms
+
 
 class BEMEngine(object):
     def __init__(self, config) -> None: