Merge ert storage

setup.py

@@ -116,7 +116,6 @@ def package_files(directory):
         "deprecation",
         "dnspython >= 2",
         "ecl >= 2.14.1",
-        "ert-storage >= 0.3.16",
         "fastapi < 0.100.0",
         "filelock",
         "graphlib_backport; python_version < '3.9'",
@@ -135,6 +134,8 @@ def package_files(directory):
         "PyQt5",
         "pyrsistent",
         "python-dateutil",
+        "python-multipart",
+        "pyarrow",
         "pyyaml",
         "qtpy",
         "requests",

src/ert/__main__.py

@@ -43,12 +43,7 @@
 
 
 def run_ert_storage(args: Namespace, _: Optional[ErtPluginManager] = None) -> None:
-    kwargs = {"ert_config": args.config, "verbose": True}
-
-    if args.database_url is not None:
-        kwargs["database_url"] = args.database_url
-
-    with StorageService.start_server(**kwargs) as server:
+    with StorageService.start_server(ert_config=args.config, verbose=True) as server:
         server.wait()
 
 
@@ -62,20 +57,13 @@ def run_webviz_ert(args: Namespace, _: Optional[ErtPluginManager] = None) -> Non
         ) from err
 
     kwargs: Dict[str, Any] = {"verbose": args.verbose}
-    if args.config:
-        ert_config = ErtConfig.from_file(args.config)
-        os.chdir(ert_config.config_path)
-        ens_path = ert_config.ens_path
-
-        # Changing current working directory means we need to
-        # only use the base name of the config file path
-        kwargs["ert_config"] = os.path.basename(args.config)
-        kwargs["project"] = os.path.abspath(ens_path)
-
-    if args.database_url is not None:
-        kwargs["database_url"] = args.database_url
+    ert_config = ErtConfig.from_file(args.config)
+    os.chdir(ert_config.config_path)
+    ens_path = ert_config.ens_path
 
-    with StorageService.init_service(**kwargs) as storage:
+    with StorageService.init_service(
+        ert_config=os.path.basename(args.config), project=os.path.abspath(ens_path)
+    ) as storage:
         storage.wait_until_ready()
         print(
             """
@@ -512,7 +500,7 @@ def ert_parser(parser: Optional[ArgumentParser], args: Sequence[str]) -> Namespa
 
 @contextmanager
 def start_ert_server(mode: str) -> Generator[None, None, None]:
-    if mode in ("api", "vis") or not FeatureToggling.is_enabled("new-storage"):
+    if mode in ("api", "vis"):
         yield
         return
 

src/ert/dark_storage/app.py

@@ -1,21 +1,49 @@
-from ert_storage.app import JSONResponse
-from ert_storage.app import app as ert_storage_app
-from ert_storage.exceptions import ErtStorageError
-from fastapi import FastAPI, Request, status
-from fastapi.openapi.docs import get_redoc_html, get_swagger_ui_html
-from fastapi.responses import HTMLResponse, RedirectResponse
+import json
+from enum import Enum
+from typing import Any
+
+from fastapi import FastAPI, Request, Response, status
+from fastapi.responses import RedirectResponse
 
 from ert.dark_storage.endpoints import router as endpoints_router
+from ert.dark_storage.exceptions import ErtStorageError
+from ert.shared import __version__
+
+
+class JSONEncoder(json.JSONEncoder):
+    """
+    Custom JSON encoder with support for Python 3.4 enums
+    """
+
+    def default(self, o: Any) -> Any:
+        if isinstance(o, Enum):
+            return o.name
+        return super().default(o)
+
+
+class JSONResponse(Response):
+    """A replacement for Starlette's JSONResponse that permits NaNs."""
+
+    media_type = "application/json"
+
+    def render(self, content: Any) -> bytes:
+        return (
+            JSONEncoder(
+                ensure_ascii=False,
+                allow_nan=True,
+                indent=None,
+                separators=(",", ":"),
+            )
+            .encode(content)
+            .encode("utf-8")
+        )
+
 
 app = FastAPI(
-    title=ert_storage_app.title,
-    version=ert_storage_app.version,
+    title="ERT Storage API (dark storage)",
+    version=__version__,
     debug=True,
     default_response_class=JSONResponse,
-    # Disable documentation so we can replace it with ERT Storage's later
-    openapi_url=None,
-    docs_url=None,
-    redoc_url=None,
 )
 
 
@@ -51,23 +79,6 @@ async def not_implemented_handler(
     return JSONResponse({}, status_code=status.HTTP_501_NOT_IMPLEMENTED)
 
 
-@app.get("/openapi.json", include_in_schema=False)
-async def get_openapi() -> JSONResponse:
-    return JSONResponse(ert_storage_app.openapi())
-
-
-@app.get("/docs", include_in_schema=False)
-async def get_swagger(req: Request) -> HTMLResponse:
-    return get_swagger_ui_html(
-        openapi_url="/openapi.json", title=f"{app.title} - Swagger UI"
-    )
-
-
-@app.get("/redoc", include_in_schema=False)
-async def get_redoc(req: Request) -> HTMLResponse:
-    return get_redoc_html(openapi_url="/openapi.json", title=f"{app.title} - Redoc")
-
-
 @app.get("/")
 async def root() -> RedirectResponse:
     return RedirectResponse("/docs")

src/ert/dark_storage/compute/__init__.py

@@ -0,0 +1 @@
+from .misfits import calculate_misfits_from_pandas

src/ert/dark_storage/compute/misfits.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Mapping, Sequence
+
+import numpy as np
+import pandas as pd
+
+if TYPE_CHECKING:
+    import numpy.typing as npt
+
+
+def _calculate_misfit(
+    obs_value: npt.NDArray[Any],
+    response_value: npt.NDArray[Any],
+    obs_std: npt.NDArray[Any],
+) -> Sequence[float]:
+    difference = response_value - obs_value
+    misfit = (difference / obs_std) ** 2
+    return (misfit * np.sign(difference)).tolist()
+
+
+def calculate_misfits_from_pandas(
+    reponses_dict: Mapping[int, pd.DataFrame],
+    observation: pd.DataFrame,
+    summary_misfits: bool = False,
+) -> pd.DataFrame:
+    """
+    Compute misfits from reponses_dict (real_id, values in dataframe)
+    and observation
+    """
+    misfits_dict = {}
+    for realization_index in reponses_dict:
+        misfits_dict[realization_index] = _calculate_misfit(
+            observation["values"],
+            reponses_dict[realization_index].loc[:, observation.index].values.flatten(),
+            observation["errors"],
+        )
+
+    df = pd.DataFrame(data=misfits_dict, index=observation.index)
+    if summary_misfits:
+        df = pd.DataFrame([df.abs().sum(axis=0)], columns=df.columns, index=[0])
+    return df.T

src/ert/dark_storage/endpoints/compute/misfits.py

@@ -4,12 +4,12 @@
 
 import pandas as pd
 from dateutil.parser import parse
-from ert_storage import exceptions as exc
-from ert_storage.compute import calculate_misfits_from_pandas
 from fastapi import APIRouter, Depends, status
 from fastapi.responses import Response
 
+from ert.dark_storage import exceptions as exc
 from ert.dark_storage.common import data_for_key, observations_for_obs_keys
+from ert.dark_storage.compute import calculate_misfits_from_pandas
 from ert.dark_storage.enkf import LibresFacade, get_res, get_storage
 from ert.storage import StorageReader
 

src/ert/dark_storage/endpoints/ensembles.py

@@ -1,9 +1,9 @@
 from typing import Any, Mapping
 from uuid import UUID
 
-from ert_storage import json_schema as js
 from fastapi import APIRouter, Body, Depends
 
+from ert.dark_storage import json_schema as js
 from ert.dark_storage.common import ensemble_parameter_names, get_response_names
 from ert.dark_storage.enkf import LibresFacade, get_res, get_storage
 from ert.storage import StorageAccessor

src/ert/dark_storage/endpoints/experiments.py

@@ -1,9 +1,9 @@
 from typing import Any, List, Mapping
 from uuid import UUID
 
-from ert_storage import json_schema as js
 from fastapi import APIRouter, Body, Depends
 
+from ert.dark_storage import json_schema as js
 from ert.dark_storage.enkf import LibresFacade, get_res, get_storage
 from ert.shared.storage.extraction import create_priors
 from ert.storage import StorageReader

src/ert/dark_storage/endpoints/observations.py

@@ -1,9 +1,9 @@
 from typing import Any, List, Mapping
 from uuid import UUID
 
-from ert_storage import json_schema as js
 from fastapi import APIRouter, Body, Depends
 
+from ert.dark_storage import json_schema as js
 from ert.dark_storage.enkf import LibresFacade, get_res
 from ert.shared.storage.extraction import create_observations
 

src/ert/dark_storage/endpoints/records.py

@@ -4,10 +4,10 @@
 from uuid import UUID, uuid4
 
 import pandas as pd
-from ert_storage import json_schema as js
 from fastapi import APIRouter, Body, Depends, File, Header, Request, UploadFile, status
 from fastapi.responses import Response
 
+from ert.dark_storage import json_schema as js
 from ert.dark_storage.common import (
     data_for_key,
     ensemble_parameters,

src/ert/dark_storage/endpoints/updates.py

@@ -1,8 +1,8 @@
 from uuid import UUID
 
-from ert_storage import json_schema as js
 from fastapi import APIRouter, Depends
 
+from ert.dark_storage import json_schema as js
 from ert.dark_storage.enkf import LibresFacade, get_res, reset_res
 
 router = APIRouter(tags=["ensemble"])

src/ert/dark_storage/enkf.py

@@ -3,10 +3,10 @@
 import os
 from typing import Optional
 
-from ert_storage.security import security
 from fastapi import Depends
 
 from ert.config import ErtConfig
+from ert.dark_storage.security import security
 from ert.enkf_main import EnKFMain
 from ert.libres_facade import LibresFacade
 from ert.storage import StorageReader, open_storage

src/ert/dark_storage/exceptions.py

@@ -0,0 +1,30 @@
+from typing import Any
+
+from fastapi import status
+
+
+class ErtStorageError(RuntimeError):
+    """
+    Base error class for all the rest of errors
+    """
+
+    __status_code__ = status.HTTP_200_OK
+
+    def __init__(self, message: str, **kwargs: Any):
+        super().__init__(message, kwargs)
+
+
+class NotFoundError(ErtStorageError):
+    __status_code__ = status.HTTP_404_NOT_FOUND
+
+
+class ConflictError(ErtStorageError):
+    __status_code__ = status.HTTP_409_CONFLICT
+
+
+class ExpectationError(ErtStorageError):
+    __status_code__ = status.HTTP_417_EXPECTATION_FAILED
+
+
+class UnprocessableError(ErtStorageError):
+    __status_code__ = status.HTTP_422_UNPROCESSABLE_ENTITY

src/ert/dark_storage/json_schema/__init__.py

@@ -0,0 +1,11 @@
+from .ensemble import EnsembleIn, EnsembleOut
+from .experiment import ExperimentIn, ExperimentOut
+from .observation import (
+    ObservationIn,
+    ObservationOut,
+    ObservationTransformationIn,
+    ObservationTransformationOut,
+)
+from .prior import Prior
+from .record import RecordOut
+from .update import UpdateIn, UpdateOut

src/ert/dark_storage/json_schema/ensemble.py

@@ -0,0 +1,37 @@
+from typing import Any, List, Mapping, Optional
+from uuid import UUID
+
+from pydantic import BaseModel, Field, root_validator
+
+
+class _Ensemble(BaseModel):
+    size: int
+    parameter_names: List[str]
+    response_names: List[str]
+    active_realizations: List[int] = []
+
+
+class EnsembleIn(_Ensemble):
+    update_id: Optional[UUID] = None
+    userdata: Mapping[str, Any] = {}
+
+    @root_validator
+    def _check_names_no_overlap(cls, values: Mapping[str, Any]) -> Mapping[str, Any]:
+        """
+        Verify that `parameter_names` and `response_names` don't overlap. Ie, no
+        record can be both a parameter and a response.
+        """
+        if not set(values["parameter_names"]).isdisjoint(set(values["response_names"])):
+            raise ValueError("parameters and responses cannot have a name in common")
+        return values
+
+
+class EnsembleOut(_Ensemble):
+    id: UUID
+    children: List[UUID] = Field(alias="child_ensemble_ids")
+    parent: Optional[UUID] = Field(alias="parent_ensemble_id")
+    experiment_id: Optional[UUID] = None
+    userdata: Mapping[str, Any]
+
+    class Config:
+        orm_mode = True

src/ert/dark_storage/json_schema/experiment.py

@@ -0,0 +1,24 @@
+from typing import Any, Mapping, Sequence
+from uuid import UUID
+
+from pydantic import BaseModel
+
+from .prior import Prior
+
+
+class _Experiment(BaseModel):
+    name: str
+
+
+class ExperimentIn(_Experiment):
+    priors: Mapping[str, Prior] = {}
+
+
+class ExperimentOut(_Experiment):
+    id: UUID
+    ensemble_ids: Sequence[UUID]
+    priors: Mapping[str, Mapping[str, Any]]
+    userdata: Mapping[str, Any]
+
+    class Config:
+        orm_mode = True