Merge branch 'flag_metadata' into 'master'

New flag metadata class

See merge request !292

geomagio/metadata/flag/Flag.py

+from typing import Dict, Union, List
+from datetime import timedelta
+
+import numpy as np
+from obspy import UTCDateTime
+from pydantic import BaseModel, Field, validator
+from enum import Enum
+
+
+class FlagCategory(str, Enum):
+    ARTIFICIAL_DISTURBANCE = "ARTIFICIAL_DISTURBANCE"
+    GAP = "GAP"
+    EVENT = "EVENT"
+    OTHER = "OTHER"
+
+
+class Flag(BaseModel):
+    """
+        Base class for flagging features in magnetic timeseries data.
+
+        Flag example:
+        ```
+        automatic_flag = Metadata(
+            created_by='ex_algorithm',
+            start_time=UTCDateTime('2023-01-01T03:05:10'),
+            end_time=UTCDateTime('2023-01-01T03:05:11'),
+            network='NT',
+            station='BOU',
+            channel='BEH',
+            category=MetadataCategory.FLAG,
+            comment="spike detected",
+            priority=1,
+            data_valid=False,
+            metadata= ArtificialDisturbance{
+                "description": "Spike in magnetic field strength",
+                "field_work": false,
+                "corrected": false,
+                "flag_category": ARTIFICIAL_DISTURBANCE,
+                "artificial_disturbance_type": ArtificialDisturbanceType.SPIKE,
+                "source": "Lightning",
+                "deviation": None,
+        }
+    )
+        ```
+    """
+
+    description: str = Field(..., description="Description of the flag")
+    field_work: bool = Field(..., description="Flag signaling field work")
+    corrected: int = Field(..., description="Corrected ID for processing stage")
+    flag_category: FlagCategory = "OTHER"
+
+
+class ArtificialDisturbanceType(str, Enum):
+    SPIKES = "SPIKES"
+    OFFSET = "OFFSET"
+    ARTIFICIAL_DISTURBANCES = "ARTIFICIAL_DISTURBANCES"
+
+
+class ArtificialDisturbance(Flag):
+    """
+    This class is used to flag artificial disturbances.
+
+    Artificial disturbances consist of the following types:
+
+    SPIKES = Single data points that are outliers in the timeseries.
+    OFFSET = A relatively constant shift or deviation in the baseline magnetic field.
+    ARTIFICIAL_DISTURBANCES = A catch-all for a continuous period of unwanted variations, may include multiple spikes, offsets and/or gaps.
+
+    Attributes
+    ----------
+    artificial_disturbance_type:ArtificialDisturbanceType
+        The type of artificial disturbance(s).
+    source: str
+        Source of the disturbance if known or suspected.
+    deviation: float
+       Deviation of an offset in nt.
+    spikes: np.ndarray
+        NumPy array of timestamps as UTCDateTime. Can be a single spike or many spikes.
+
+    """
+
+    artificial_disturbance_type: ArtificialDisturbanceType
+    deviation: float = None
+    source: str = None
+    spikes: np.ndarray = None
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.flag_category = "ARTIFICIAL_DISTURBANCE"
+
+
+class Gap(Flag):
+    """
+    This class is used to flag gaps in data.
+
+    A gap is a period where data is missing or not recorded.
+
+    Attributes
+    ----------
+    cause: str
+        Cause of gap, e.g., network outage.
+    handling: str
+        How the gap is being handled, e.g., backfilled.
+    """
+
+    cause: str = None
+    handling: str = None
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.flag_category = "GAP"
+
+
+class EventType(str, Enum):
+    GEOMAGNETIC_STORM = "GEOMAGNETIC_STORM"
+    GEOMAGNETIC_SUBSTORM = "GEOMAGNETIC_SUBSTORM"
+    EARTHQUAKE = "EARTHQUAKE"
+    OTHER = "OTHER"
+
+
+class Event(Flag):
+    """
+    This class is used to flag an event of interest such as a geomagnetic storm or earthquake.
+
+    Attributes
+    ----------
+    event_type : EventType
+        The type of event.
+    scale : str
+        Geomagnetic storm scale or Richter scale magnitude.
+    index : int
+        Planetary K-index, DST index or some other index.
+    url : str
+        A url related to the event. Could be NOAA SWPC, USGS Earthquakes page or another site.
+    """
+
+    event_type: EventType
+    index: int = None
+    scale: str = None
+    url: str = None
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.flag_category = "EVENT"
+
+
+# More example usage:
+timestamps_array = np.array(
+    [
+        UTCDateTime("2023-11-16T12:00:0"),
+        UTCDateTime("2023-11-16T12:01:10"),
+        UTCDateTime("2023-11-16T12:02:30"),
+    ]
+)
+
+spikes_data = {
+    "starttime": "2023-11-16 12:00:00",
+    "endtime": "2023-11-16 12:02:30",
+    "description": "Spikes description",
+    "field_work": False,
+    "corrected": 32265,
+    "disturbance_type": ArtificialDisturbanceType.SPIKES,
+    "source": "processing",
+    "spikes": timestamps_array,
+}
+
+offset_data = {
+    "description": "Offset description",
+    "field_work": False,
+    "corrected": 47999,
+    "disturbance_type": ArtificialDisturbanceType.OFFSET,
+    "source": "Bin change",
+    "deviation": 10.0,
+}
+geomagnetic_storm_data = {
+    "description": "Geomagnetic storm",
+    "field_work": False,
+    "corrected": 36999,
+    "event_type": EventType.GEOMAGNETIC_STORM,
+    "scale": "G3",
+    "index": 7,
+    "url": "https://www.swpc.noaa.gov/products/planetary-k-index",
+}
+
+spike_instance = ArtificialDisturbance(**spikes_data)
+offset_instance = ArtificialDisturbance(**offset_data)
+
+print(spike_instance.dict())
+print(offset_instance.dict())