earthscope-sdk 1.1.0__py3-none-any.whl → 1.2.0b0__py3-none-any.whl

earthscope_sdk/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.1.0"
+__version__ = "1.2.0b0"
 
 from earthscope_sdk.client import AsyncEarthScopeClient, EarthScopeClient
 

earthscope_sdk/client/_client.py

@@ -8,6 +8,26 @@ class AsyncEarthScopeClient(AsyncSdkClient):
     An async client for interacting with api.earthscope.org
     """
 
+    @cached_property
+    def data(self):
+        """
+        Data access functionality
+        """
+        # lazy load
+        from earthscope_sdk.client.data_access._service import AsyncDataAccessService
+
+        return AsyncDataAccessService(self)
+
+    @cached_property
+    def discover(self):
+        """
+        Data discovery functionality
+        """
+        # lazy load
+        from earthscope_sdk.client.discovery._service import AsyncDiscoveryService
+
+        return AsyncDiscoveryService(self._ctx)
+
     @cached_property
     def user(self):
         """
@@ -24,6 +44,33 @@ class EarthScopeClient(SdkClient):
     A client for interacting with api.earthscope.org
     """
 
+    @cached_property
+    def _async_client(self):
+        """
+        An async client for interacting with api.earthscope.org
+        """
+        return AsyncEarthScopeClient(ctx=self._ctx)
+
+    @cached_property
+    def data(self):
+        """
+        Data access functionality
+        """
+        # lazy load
+        from earthscope_sdk.client.data_access._service import DataAccessService
+
+        return DataAccessService(self._async_client)
+
+    @cached_property
+    def discover(self):
+        """
+        Data discovery functionality
+        """
+        # lazy load
+        from earthscope_sdk.client.discovery._service import DiscoveryService
+
+        return DiscoveryService(self._ctx)
+
     @cached_property
     def user(self):
         """

earthscope_sdk/client/data_access/_arrow/_common.py

@@ -0,0 +1,94 @@
+try:
+    import pyarrow as pa
+except ModuleNotFoundError as e:
+    raise ModuleNotFoundError(
+        "Optional dependency 'pyarrow' is required for this feature. "
+        "Install it with: pip install earthscope-sdk[arrow]"
+    ) from e
+
+
+from earthscope_sdk.client.discovery.models import DatasourceBaseModel
+
+
+def convert_time(
+    table: pa.Table,
+    *,
+    source_field: str = "time",
+    target_field: str = "timestamp",
+    unit: str = "ms",
+    tz: str = "UTC",
+) -> pa.Table:
+    """
+    Convert a unix timestamp (integer since epoch) field to a timestamp type.
+
+    Args:
+        table: The table to convert.
+        source_field: The name of the field to convert.
+        target_field: The name of the field to store the result.
+        unit: The unit of the time field.
+        tz: The timezone of the time field.
+
+    Returns:
+        The converted table.
+    """
+    if unit not in {"s", "ms", "us", "ns"}:
+        raise ValueError(f"Invalid time unit: {unit}")
+
+    time_idx = table.schema.get_field_index(source_field)
+    if time_idx < 0:
+        return table
+
+    return table.set_column(
+        time_idx,
+        target_field,
+        table[source_field].cast(pa.timestamp(unit, tz)),
+    )
+
+
+def get_datasource_metadata_table(
+    datasources: list[DatasourceBaseModel],
+    *,
+    fields: list[str],
+    namespaces: list[str],
+) -> pa.Table:
+    """
+    Get a pyarrow table containing metadata columns for a list of datasources.
+
+    Args:
+        datasources: The list of datasources to get metadata for.
+        fields: The fields to include in the metadata table.
+        namespaces: The namespaces to include in the metadata table.
+
+    Returns:
+        The metadata table.
+    """
+    return pa.Table.from_pylist(
+        [s.to_arrow_columns(fields=fields, namespaces=namespaces) for s in datasources]
+    )
+
+
+def load_table_with_extra(content: bytes, *, is_stream=True, **kwargs) -> pa.Table:
+    """
+    Load a table from an arrow stream or file. Optionally add extra metadata columns.
+
+    Args:
+        content: The content of the stream or file.
+        is_stream: Whether the content is an arrow stream or an arrow file.
+        **kwargs: The extra metadata columns to add to the table.
+
+    Returns:
+        The loaded table.
+    """
+    if is_stream:
+        reader = pa.ipc.open_stream(content)
+    else:
+        reader = pa.ipc.open_file(content)
+
+    table = reader.read_all()
+
+    if kwargs:
+        count = len(table)
+        for key, value in kwargs.items():
+            table = table.append_column(key, pa.repeat(value, count))
+
+    return table

earthscope_sdk/client/data_access/_arrow/_gnss.py

@@ -0,0 +1,116 @@
+import pyarrow as pa
+import pyarrow.compute as pc
+
+from earthscope_sdk.client.data_access._arrow._common import convert_time
+
+_SYSTEM_ID = "system_id"
+_SYSTEM_TABLE = pa.Table.from_pylist(
+    [
+        {_SYSTEM_ID: 1, "system": "G"},  # GPS
+        {_SYSTEM_ID: 2, "system": "R"},  # GLONASS
+        {_SYSTEM_ID: 3, "system": "S"},  # SBAS
+        {_SYSTEM_ID: 4, "system": "E"},  # Galileo
+        {_SYSTEM_ID: 5, "system": "C"},  # BeiDou
+        {_SYSTEM_ID: 6, "system": "J"},  # QZSS
+        {_SYSTEM_ID: 7, "system": "I"},  # IRNSS/NavIC
+    ],
+    schema=pa.schema(
+        [
+            pa.field(_SYSTEM_ID, pa.uint8()),
+            pa.field("system", pa.string()),
+        ]
+    ),
+)
+
+
+def _convert_obs_code(
+    table: pa.Table,
+    *,
+    source_field: str = "obs",
+    target_field: str = "obs_code",
+) -> pa.Table:
+    """
+    Convert the obs_code field from packed chars to a string.
+
+    Args:
+        table: The table to convert.
+        source_field: The name of the field to convert.
+        target_field: The name of the field to store the result.
+
+    Returns:
+        The converted table.
+    """
+    # extract obs_code (encoded as packed chars)
+    result_chunks = [
+        pc.cast(
+            pa.FixedSizeBinaryArray.from_buffers(
+                type=pa.binary(2),
+                length=len(chunk),
+                buffers=[None, chunk.buffers()[1]],
+            ),
+            pa.string(),
+        )
+        for chunk in table[source_field].chunks
+    ]
+
+    return table.set_column(
+        table.schema.get_field_index(source_field),
+        target_field,
+        # can't easily swap endianness of packed chars, so reverse the string
+        pc.utf8_reverse(pa.chunked_array(result_chunks, type=pa.string())),
+    )
+
+
+def _convert_system(table: pa.Table, *, field: str) -> pa.Table:
+    """
+    Convert the sys field from integer to string.
+
+    Args:
+        table: The table to convert.
+        field: The name of the field to convert.
+
+    Returns:
+        The converted table.
+    """
+    return table.join(_SYSTEM_TABLE, field, right_keys=_SYSTEM_ID).drop(field)
+
+
+def prettify_observations_table(table: pa.Table) -> pa.Table:
+    """
+    Clean up a raw observations table (as stored in TileDB) for easier use.
+
+    Args:
+        table: The table to prettify.
+
+    Returns:
+        The prettified table.
+    """
+    # Cast time column (integer milliseconds since epoch) to timestamp type
+    table = convert_time(table)
+
+    # Expose friendlier column names
+    table = table.rename_columns({"sat": "satellite"})
+
+    # Convert system column from integer to string
+    table = _convert_system(table, field="sys")
+
+    # Convert obs_code column from packed chars to string
+    table = _convert_obs_code(table)
+
+    return table
+
+
+def prettify_positions_table(table: pa.Table) -> pa.Table:
+    """
+    Clean up a raw positions table (as stored in TileDB) for easier use.
+
+    Args:
+        table: The table to prettify.
+
+    Returns:
+        The prettified table.
+    """
+    # Convert time column (integer milliseconds since epoch) to timestamp type
+    table = convert_time(table)
+
+    return table

earthscope_sdk/client/data_access/_base.py

@@ -0,0 +1,85 @@
+import datetime as dt
+from typing import TYPE_CHECKING, Union
+
+from earthscope_sdk.client.data_access._arrow._common import load_table_with_extra
+from earthscope_sdk.common.service import SdkService
+from earthscope_sdk.util._time import to_utc_dt
+
+if TYPE_CHECKING:
+    from pyarrow import Table
+
+
+class DataAccessBaseService(SdkService):
+    """
+    L1 data access service functionality
+    """
+
+    async def _get_gnss_observations(
+        self,
+        *,
+        session_edid: str,
+        start_datetime: dt.datetime,
+        end_datetime: dt.datetime,
+        system: Union[str, list[str]] = [],
+        satellite: Union[str, list[str]] = [],
+        obs_code: Union[str, list[str]] = [],
+        field: Union[str, list[str]] = [],
+    ) -> "Table":
+        """
+        Retrieve GNSS observations.
+
+        This method retrieves GNSS observations for a given session.
+
+        Args:
+            session_edid: The session EDID to retrieve observations for.
+            start_datetime: The start datetime to retrieve positions for.
+            end_datetime: The end datetime to retrieve positions for.
+            system: The system(s) to retrieve observations for.
+            satellite: The satellite(s) to retrieve observations for.
+            obs_code: The observation code(s) to retrieve observations for.
+            field: The field(s) to retrieve observations for.
+
+        Returns:
+            A pyarrow table containing the GNSS observations.
+        """
+
+        # lazy import
+
+        headers = {
+            "accept": "application/vnd.apache.arrow.stream",
+        }
+
+        params = {
+            "session_id": session_edid,
+            "start_datetime": to_utc_dt(start_datetime).isoformat(),
+            "end_datetime": to_utc_dt(end_datetime).isoformat(),
+        }
+
+        if system:
+            params["system"] = system
+
+        if satellite:
+            params["satellite"] = satellite
+
+        if obs_code:
+            params["obs_code"] = obs_code
+
+        if field:
+            params["field"] = field
+
+        req = self.ctx.httpx_client.build_request(
+            method="GET",
+            url=f"{self.resources.api_url}beta/data-products/gnss/observations",
+            headers=headers,
+            params=params,
+            timeout=30,
+        )
+
+        resp = await self._send_with_retries(req)
+
+        return await self.ctx.run_in_executor(
+            load_table_with_extra,
+            resp.content,
+            # Add edid column to the table returned by the API
+            edid=session_edid,
+        )

earthscope_sdk/client/data_access/_query_plan/_gnss_observations.py

@@ -0,0 +1,295 @@
+import datetime as dt
+from typing import TYPE_CHECKING, Literal, NamedTuple, Optional
+
+from typing_extensions import Self
+
+from earthscope_sdk.client.data_access._arrow._common import (
+    get_datasource_metadata_table,
+)
+from earthscope_sdk.client.data_access._arrow._gnss import prettify_observations_table
+from earthscope_sdk.client.data_access._query_plan._query_plan import (
+    AsyncQueryPlan,
+    QueryPlan,
+)
+from earthscope_sdk.client.discovery.models import SessionDatasource
+from earthscope_sdk.util._itertools import to_set
+from earthscope_sdk.util._time import TimePeriod, time_range_periods
+from earthscope_sdk.util._types import ListOrItem
+
+if TYPE_CHECKING:
+    import pyarrow as pa
+
+    from earthscope_sdk.client import AsyncEarthScopeClient
+
+
+_MAX_EPOCHS_PER_REQUEST = 11520
+"""
+Maximum number of epochs per GNSS observations request.
+"""
+
+
+_NAMESPACES = {"igs", "4charid", "dataflow"}
+
+
+GnssObservationField = Literal[
+    "range",
+    "phase",
+    "doppler",
+    "snr",
+    "slip",
+    "flags",
+    "fcn",
+]
+"""
+Fields available when fetching GNSS observations.
+
+- range: code / psuedorange
+- phase: carrier phase
+- doppler: doppler shift
+- snr: signal to noise ratio
+- slip: carrier phase cycle slip occurred
+- flags: event flags
+- fcn: GLONASS frequency channel number
+"""
+
+
+GnssObservationsMetaField = Literal[
+    "edid",
+    "igs",
+    "4charid",
+    "dataflow",
+    "sample_interval",
+    "roll",
+]
+"""
+Metadata fields available for joining to GNSS observations.
+"""
+
+SatelliteSystem = Literal["G", "R", "E", "J", "C", "I", "S"]
+"""
+GNSS satellite system abbreviation.
+
+- G: GPS
+- R: GLONASS
+- E: Galileo
+- J: QZSS
+- C: BeiDou
+- I: IRNSS / NavIC
+- S: SBAS
+"""
+
+
+class GnssObservationsRequest(NamedTuple):
+    """
+    An individual request for GNSS observations.
+    """
+
+    period: TimePeriod
+    session: SessionDatasource
+
+
+def _get_max_request_period(session: SessionDatasource) -> dt.timedelta:
+    """
+    Get the maximum request period for a session.
+
+    Even though the API supports up to 115200 epochs per request, we limit
+    request periods to nice round numbers that easily group into days.
+    """
+
+    if session.sample_interval == dt.timedelta(milliseconds=200):
+        return dt.timedelta(minutes=30)
+
+    if session.sample_interval == dt.timedelta(seconds=1):
+        return dt.timedelta(hours=3)
+
+    # fallback to actual limit of 11520 epochs per request
+    # 15s -> 2 days
+    # 30s -> 4 days
+    # 1m -> 8 days
+    # etc.
+    return session.sample_interval * _MAX_EPOCHS_PER_REQUEST
+
+
+class AsyncGnssObservationsQueryPlan(
+    AsyncQueryPlan[GnssObservationsRequest],
+):
+    """
+    A query plan for GNSS observations.
+    """
+
+    def __init__(
+        self,
+        client: "AsyncEarthScopeClient",
+        # observation parameters
+        start_datetime: dt.datetime,
+        end_datetime: dt.datetime,
+        system: ListOrItem[SatelliteSystem] = [],
+        satellite: ListOrItem[str] = [],
+        obs_code: ListOrItem[str] = [],
+        field: ListOrItem[GnssObservationField] = [],
+        # session parameters
+        network_name: ListOrItem[str] = [],
+        network_edid: ListOrItem[str] = [],
+        station_name: ListOrItem[str] = [],
+        station_edid: ListOrItem[str] = [],
+        session_name: ListOrItem[str] = [],
+        session_edid: ListOrItem[str] = [],
+        roll: Optional[dt.timedelta] = None,
+        sample_interval: Optional[dt.timedelta] = None,
+        # metadata parameters
+        meta_fields: ListOrItem[GnssObservationsMetaField] = ["igs"],
+    ):
+        if (
+            not network_name
+            and not network_edid
+            and not station_name
+            and not station_edid
+            and not session_name
+            and not session_edid
+        ):
+            raise ValueError("Expected at least one station or session name or edid")
+
+        super().__init__(client)
+
+        # Common parameters
+        self.system = system
+        self.satellite = satellite
+        self.obs_code = obs_code
+        self.field = field
+
+        # Request parameters
+        self.period = TimePeriod(start=start_datetime, end=end_datetime)
+
+        # Session parameters
+        self.network_name = network_name
+        self.network_edid = network_edid
+        self.station_name = station_name
+        self.station_edid = station_edid
+        self.session_name = session_name
+        self.session_edid = session_edid
+        self.roll = roll
+        self.sample_interval = sample_interval
+
+        # Metadata parameters
+        meta_fields = to_set(meta_fields)
+        self.meta_namespaces = list(meta_fields & _NAMESPACES)
+        self.meta_fields = list(meta_fields - _NAMESPACES)
+
+        # Local state
+        self._meta_table: Optional["pa.Table"] = None
+        self._max_request_period: Optional[dt.timedelta] = None
+
+    async def _build_requests(self) -> list[GnssObservationsRequest]:
+        # Expand station/session names into session_edids
+        sessions: list[SessionDatasource]
+        sessions = await self._client.discover.list_session_datasources(
+            network_name=self.network_name,
+            network_edid=self.network_edid,
+            station_name=self.station_name,
+            station_edid=self.station_edid,
+            session_name=self.session_name,
+            session_edid=self.session_edid,
+            roll=self.roll,
+            sample_interval=self.sample_interval,
+            with_parents=True,
+        )
+
+        # Pre-compute metadata table
+        if self.meta_fields or self.meta_namespaces:
+            self._meta_table = get_datasource_metadata_table(
+                sessions,
+                fields=self.meta_fields + ["edid"],  # edid is required for join
+                namespaces=self.meta_namespaces,
+            )
+
+        requests: list[GnssObservationsRequest] = []
+        for session in sessions:
+            max_period = _get_max_request_period(session)
+            if self._max_request_period is not None:
+                max_period = min(max_period, self._max_request_period)
+
+            time_periods = time_range_periods(
+                start=self.period.start,
+                end=self.period.end,
+                period=max_period,
+            )
+            requests.extend(GnssObservationsRequest(p, session) for p in time_periods)
+
+        return requests
+
+    async def _execute_one(self, args: GnssObservationsRequest) -> Optional["pa.Table"]:
+        from httpx import HTTPStatusError  # lazy import
+
+        try:
+            return await self._client.data._get_gnss_observations(
+                start_datetime=args.period.start,
+                end_datetime=args.period.end,
+                session_edid=args.session.edid,
+                system=self.system,
+                satellite=self.satellite,
+                obs_code=self.obs_code,
+                field=self.field,
+            )
+        except HTTPStatusError as e:
+            if e.response.status_code != 404:
+                raise e
+
+            return None
+
+    def _hook(self, table: "pa.Table") -> "pa.Table":
+        table = prettify_observations_table(table)
+
+        # Add metadata columns
+        if self._meta_table:
+            table = table.join(self._meta_table, "edid")
+
+        # Drop edid column if not requested
+        if "edid" not in self.meta_fields:
+            table = table.drop_columns(["edid"])
+
+        return table
+
+    def group_by_day(self) -> Self:
+        """
+        Group the requests by day.
+
+        This will configure the query plan to fetch data for all stations
+        one day at a time.
+
+        NOTE: This groups requests by their start times. Should a request span a
+        day boundary, it will only be included in the day of its start time.
+        """
+        # Truncate requests to the day of their start time
+        self._max_request_period = dt.timedelta(days=1)
+
+        return self.group_by(lambda r: r.period.start.date())
+
+    def group_by_station(self) -> Self:
+        """
+        Group the requests by station.
+
+        This will configure the query plan to fetch data for the entire time
+        range one station at a time.
+        """
+        return self.group_by(lambda r: r.session.station.edid)
+
+    def sort_by_station(self, *, reverse: bool = False) -> Self:
+        """
+        Sort the requests by station name.
+        """
+        return self.sort_by(lambda r: r.session.station.names["IGS"], reverse=reverse)
+
+
+class GnssObservationsQueryPlan(
+    QueryPlan[GnssObservationsRequest],
+):
+    """
+    A query plan for GNSS observations.
+    """
+
+    def __init__(self, async_plan: AsyncGnssObservationsQueryPlan):
+        super().__init__(async_plan)
+
+        self.group_by_day = self._wrap_and_return_self(async_plan.group_by_day)
+        self.group_by_station = self._wrap_and_return_self(async_plan.group_by_station)
+        self.sort_by_station = self._wrap_and_return_self(async_plan.sort_by_station)