Support for SMRY, EGRID, INIT and restarts via resdata (#14)

.gitignore

@@ -2,3 +2,4 @@
 /docs/Manifest.toml
 /docs/build/
 .vscode
+.CondaPkg/*

CondaPkg.toml

@@ -0,0 +1,6 @@
+[deps]
+python = ""
+
+[pip.deps]
+resdata = ""
+numpy = ""

Project.toml

@@ -1,7 +1,7 @@
 name = "GeoEnergyIO"
 uuid = "3b1dd628-313a-45bb-9d8d-8f3c48dcb5d4"
 authors = ["Olav Møyner <olav.moyner@gmail.com> and contributors"]
-version = "1.1.19"
+version = "1.1.20"
 
 [deps]
 Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
@@ -15,6 +15,12 @@ Parsers = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
+[weakdeps]
+PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
+
+[extensions]
+GeoEnergyIOPythonCallExt = "PythonCall"
+
 [compat]
 Artifacts = "1"
 DelimitedFiles = "1.6.7"
@@ -29,6 +35,7 @@ julia = "1.8"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 
 [targets]
 test = ["Test"]

docs/src/index.md

@@ -149,3 +149,32 @@ GeoEnergyIO.InputParser.keyword_default_value
 
 ```@index
 ```
+
+# Extensions
+
+## Simulation outputs supported via resdata package
+
+The Python package [resdata](https://github.com/equinor/resdata) developed by Equinor can be loaded to add support for reading summary files (sparse data), egrid (processed grid), init (initial conditions) and restart (cell-wise results).
+
+To add support for this extension, you have to add `PythonCall` to your environment (one-time operation):
+
+```julia
+using Pkg
+Pkg.add("PythonCall")
+```
+
+Afterwards, you can then load the package to get access to the new functions:
+
+```julia
+using PythonCall
+```
+
+```@docs
+read_restart
+read_init
+read_egrid
+read_summary
+```
+
+!!! note "resdata is GPL-3.0 licensed"
+    The resdata package is under a different license than GeoEnergyIO which uses MIT. The licenses are compatible, but a distributed product that contains resdata must comply with the terms of the GPL license.

ext/GeoEnergyIOPythonCallExt/GeoEnergyIOPythonCallExt.jl

@@ -0,0 +1,162 @@
+module GeoEnergyIOPythonCallExt
+    import GeoEnergyIO
+    using PythonCall
+
+    function GeoEnergyIO.read_summary_impl(pth; extra_out = false, verbose = false)
+        summary_mod = pyimport("resdata.summary")
+        smry = summary_mod.Summary(pth);
+        smry_keys = to_julia_keys(smry)
+
+        t_total = 0.0
+        out = Dict{String, Any}()
+        for k in smry_keys
+            t_kw = @elapsed kw = pyconvert(Vector{Float64}, smry.numpy_vector(k))
+            if ':' in k
+                cat, name = split(k, ':')
+                if haskey(out, cat)
+                    out[cat][name] = kw
+                else
+                    out[cat] = Dict{String, Any}(name => kw)
+                end
+            else
+                out[k] = kw
+            end
+            if verbose
+                println("$k read in $t_kw s")
+            end
+            t_total += t_kw
+        end
+        if extra_out
+            ret = (out, smry)
+        else
+            ret = out
+        end
+        if verbose
+            println("resdata extension: Read $(length(smry_keys)) keywords in $t_total s")
+        end
+        return ret
+    end
+
+    function GeoEnergyIO.read_egrid_impl(pth; extra_out = false, verbose = false)
+        grid_mod = pyimport("resdata.grid")
+        epth = "$pth.EGRID"
+        isfile(epth) || error("File not found: $epth")
+        egrid = grid_mod.Grid(epth)
+
+        nx, ny, nz, nact = pyconvert(Tuple, egrid.getDims())
+        out = Dict{String, Any}()
+        actnum = to_julia_array(egrid.export_actnum(), Int64)
+        @assert length(actnum) == nx*ny*nz
+        actnum_actual = Array{Bool, 3}(undef, nx, ny, nz)
+        for i in eachindex(actnum, actnum_actual)
+            actnum_actual[i] = actnum[i] > 0
+        end
+        out["ACTNUM"] = actnum_actual
+        out["ZCORN"] = to_julia_array(egrid.export_zcorn())
+        out["COORD"] = to_julia_array(egrid.export_coord())
+        out["cartDims"] = [nx, ny, nz]
+
+        if extra_out
+            ret = (out, egrid)
+        else
+            ret = out
+        end
+        return ret
+    end
+
+    function GeoEnergyIO.read_restart_impl(pth; extra_out = false, actnum = missing, egrid = missing)
+        resfile_mod = pyimport("resdata.resfile")
+        grid_mod = pyimport("resdata.grid")
+        np = pyimport("numpy")
+
+        basepth, ext = splitext(pth)
+        @info basepth
+        if ismissing(egrid)
+            egrid = grid_mod.Grid("$basepth.EGRID")
+        end
+        if ext == ""
+            pth = "$pth.UNRST"
+        end
+        rstrt = resfile_mod.ResdataRestartFile(egrid, filename = pth)
+        hkeys = to_julia_keys(rstrt)
+        out = Dict{String, Any}[]
+        for h in rstrt.headers()
+            step = Dict{String, Any}()
+            step["days"] = pyconvert(Float64, h.get_sim_days())
+            step["date"] = pyconvert(String, h.get_sim_date().strftime("%Y-%m-%dT%H:%M:%S"))
+            step["report_step"] = pyconvert(Int64, h.get_report_step())
+            rst_block = rstrt.restart_view(report_step = h.get_report_step())
+            for k in hkeys
+                v = rst_block[k][0]
+                dtype = v.dtype
+                if pyconvert(Bool, dtype == np.int32)
+                    T = Int64
+                elseif pyconvert(Bool, dtype == np.float64) || pyconvert(Bool, dtype == np.float32)
+                    T = Float64
+                else
+                    continue
+                end
+                v_array = to_julia_array(v)
+                step[k] = map_to_active(v_array, actnum)
+            end
+            push!(out, step)
+        end
+        if extra_out
+            ret = (out, rstrt)
+        else
+            ret = out
+        end
+        return ret
+    end
+
+    function GeoEnergyIO.read_init_impl(pth; extra_out = false, actnum = missing)
+        np = pyimport("numpy")
+        resfile_mod = pyimport("resdata.resfile")
+        ipth = "$pth.INIT"
+        isfile(ipth) || error("File not found: $ipth")
+        init = resfile_mod.ResdataFile(ipth)
+        out = Dict{String, Any}()
+        for k in to_julia_keys(init)
+            v = init[k][0]
+            dtype = v.dtype
+            if pyconvert(Bool, dtype == np.int32)
+                T = Int64
+            elseif pyconvert(Bool, dtype == np.float64) || pyconvert(Bool, dtype == np.float32)
+                T = Float64
+            else
+                continue
+            end
+            ju_vec = to_julia_array(v, T)
+            out[k] = map_to_active(ju_vec, actnum)
+        end
+        if extra_out
+            ret = (out, init)
+        else
+            ret = out
+        end
+        return ret
+    end
+
+    function to_julia_array(pyarr, T = Float64)
+        return pyconvert(Vector{T}, pyarr.numpy_copy())
+    end
+
+    function to_julia_keys(x)
+        return pyconvert(Vector{String}, collect(x.keys()))
+    end
+
+    function map_to_active(ju_vec, actnum::AbstractArray)
+        return map_to_active(ju_vec, vec(actnum))
+    end
+
+    function map_to_active(ju_vec, actnum::Vector)
+        if length(ju_vec) == length(actnum)
+            ju_vec = ju_vec[actnum]
+        end
+        return ju_vec
+    end
+
+    function map_to_active(ju_vec, ::Missing)
+        return ju_vec
+    end
+end

src/GeoEnergyIO.jl

@@ -13,6 +13,9 @@ module GeoEnergyIO
     include("CornerPointGrid/CornerPointGrid.jl")
     import .CornerPointGrid: mesh_from_grid_section
 
+    export read_restart, read_init, read_egrid, read_summary
+    include("ext.jl")
+
     function test_input_file_path(dataset::AbstractString, filename = missing)
         pth = @artifact_str(dataset)
         if !ismissing(filename)

src/ext.jl

@@ -0,0 +1,120 @@
+function check_pythoncall(; throw = true)
+    present = !isnothing(Base.get_extension(GeoEnergyIO, :GeoEnergyIOPythonCallExt))
+    if throw
+        if !present
+            error("PythonCall is not loaded. This is needed for reading of output files via resdata.\nTo fix:\nusing Pkg; Pkg.add(\"PythonCall\"); using PythonCall")
+        end
+    end
+    return present
+end
+
+"""
+    restart = read_restart(fn)
+    restart, raw_restart = read_restart(fn, extra_out = true)
+
+Read a restart file from `fn`. This should be the base path (i.e. without the
+`.RSRT` extension). The results are given as a Vector of Dicts.
+
+# Keyword arguments
+
+- `extra_out`: If true, return the raw Python object as well as the parsed
+  data. Default is false.
+- `actnum=missing`: ACTNUM array that can be used to reduce the outputs to the
+  active cells.
+- `egrid=missing`: EGRID object needed to read the restarts. Will be read from
+  the same path as `fn` if not provided.
+
+# Notes
+
+This function requires the `resdata` Python package to be installed, which will
+be automatically added to your environment if you first install PythonCall and
+put `using PythonCall` in your script or REPL.
+
+The main class to lookup on the Python side of things is `ResdataRestartFile`.
+"""
+function read_restart(arg...; kwarg...)
+    check_pythoncall()
+    return read_restart_impl(arg...; kwarg...)
+end
+
+function read_restart_impl
+
+end
+
+"""
+    init = read_init(fn)
+    init, raw_init = read_init(fn, extra_out = true)
+
+Read a init file from `fn`. This should be the base path (i.e. without the
+`.RSRT` extension). The results are given as a Dict.
+
+# Keyword arguments
+
+- `extra_out`: If true, return the raw Python object as well as the parsed
+  data. Default is false.
+- `actnum=missing`: ACTNUM array that can be used to reduce the outputs to the
+  active cells.
+
+# Notes
+
+This function requires the `resdata` Python package to be installed, which will
+be automatically added to your environment if you first install PythonCall and
+put `using PythonCall` in your script or REPL.
+
+The main class to lookup on the Python side of things is `ResdataFile`.
+"""
+function read_init(arg...; kwarg...)
+    check_pythoncall()
+    return read_init_impl(arg...; kwarg...)
+end
+
+function read_init_impl
+
+end
+
+"""
+    egrid = read_egrid(pth)
+    egrid, raw_egrid = read_egrid(pth, extra_out = true)
+
+Read the EGRID file from `pth`. The results are given as a Dict and can be
+passed further on to `mesh_from_grid_section` to construct a Jutul mesh.
+
+# Notes
+
+This function requires the `resdata` Python package to be installed, which will
+be automatically added to your environment if you first install PythonCall and
+put `using PythonCall` in your script or REPL.
+
+Uses primarily `resdata.grid.Grid`.
+"""
+function read_egrid(arg...; kwarg...)
+    check_pythoncall()
+    return read_egrid_impl(arg...; kwarg...)
+end
+
+function read_egrid_impl
+
+end
+
+"""
+    summary = read_summary(pth)
+    summary, raw_summary = read_summary(pth, extra_out = true)
+
+Read the SUMMARY file from `pth`. The results are given as a Dict.
+
+# Notes
+
+This function requires the `resdata` Python package to be installed, which will
+be automatically added to your environment if you first install PythonCall and
+put `using PythonCall` in your script or REPL.
+
+Uses primarily `resdata.summary.Summary`.
+"""
+function read_summary(arg...; kwarg...)
+    check_pythoncall()
+    return read_summary_impl(arg...; kwarg...)
+end
+
+function read_summary_impl
+
+end