feat: overwrite DataFrame.from_dict to make sure fields have the righ…

…t dtype
swiss-seismological-service · Sep 28, 2023 · 6d16259 · 6d16259
1 parent de055f4
commit 6d16259
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 19 deletions.
diff --git a/catalog_tools/io/client.py b/catalog_tools/io/client.py
@@ -88,18 +88,4 @@ def get_events(self, start_time: datetime | None = None,
         r = requests.get(request_url, stream=True)
         catalog = parse_quakeml_response(r)
 
-        df = Catalog.from_dict(catalog)
-
-        if not include_uncertainty:
-            df = df.drop_uncertainties()
-
-        if not include_ids:
-            df = df.drop_ids()
-
-        df['magnitude'] = pd.to_numeric(df['magnitude'])
-        df['latitude'] = pd.to_numeric(df['latitude'])
-        df['longitude'] = pd.to_numeric(df['longitude'])
-        df['depth'] = pd.to_numeric(df['depth'])
-        df['time'] = pd.to_datetime(df['time'])
-
-        return df
+        return Catalog.from_dict(catalog, include_uncertainty, include_ids)
diff --git a/catalog_tools/seismicity/catalog.py b/catalog_tools/seismicity/catalog.py
@@ -75,10 +75,43 @@ def from_quakeml(cls, quakeml: str,
         else:
             catalog = parse_quakeml(quakeml, includeallmagnitudes)
 
-        df = cls.from_dict(catalog)
+        df = cls.from_dict(catalog, includeuncertainties, includeids)
 
-        if not includeuncertainties:
+        return df
+
+    @classmethod
+    def from_dict(cls,
+                  data: list[dict],
+                  includeuncertainty: bool = True,
+                  includeids: bool = True, *args, **kwargs) -> Catalog:
+        """
+        Create a Catalog from a list of dictionaries.
+
+        Args:
+            data : list[dict]
+                A list of earthquake event information dictionaries.
+
+        Returns:
+            Catalog
+        """
+
+        df = super().from_dict(data, *args, **kwargs)
+        df = cls(df)
+
+        if 'magnitude' in df.columns:
+            df['magnitude'] = pd.to_numeric(df['magnitude'])
+        if 'latitude' in df.columns:
+            df['latitude'] = pd.to_numeric(df['latitude'])
+        if 'longitude' in df.columns:
+            df['longitude'] = pd.to_numeric(df['longitude'])
+        if 'depth' in df.columns:
+            df['depth'] = pd.to_numeric(df['depth'])
+        if 'time' in df.columns:
+            df['time'] = pd.to_datetime(df['time']).dt.tz_localize(None)
+
+        if not includeuncertainty:
             df = df.drop_uncertainties()
+
         if not includeids:
             df = df.drop_ids()
 
@@ -200,7 +233,7 @@ def _create_ids(self):
 
         return self
 
-    @require_cols(require=_required_cols)
+    @require_cols(require=_required_cols + ['magnitude_type'])
     def to_quakeml(self, agencyID=' ', author=' ') -> str:
         """
         Convert the catalog to QuakeML format.

diff --git a/catalog_tools/seismicity/catalog_templates/quakeml.j2 b/catalog_tools/seismicity/catalog_templates/quakeml.j2
@@ -38,7 +38,7 @@
 
             <origin publicID="{{ event.originid }}">
                 <time>
-                    <value>{{ event.time }}</value>
+                    <value>{{ event.time.strftime('%Y-%m-%dT%H:%M:%S.%fZ') }}</value>
                 </time>
                 <longitude>
                     <value>{{ event.longitude }}</value>