A lift table from NeqSim

"""A lift table NeqSim makes: a producer's tubing and its oil's components in, a table out.

``lift_table.py`` builds its table by hand. Here NeqSim makes it: PetroScript describes the
producer's tubing (``ps.lift.Tubing``) and the stream it flows, an oil described by its components
(``ps.fluid.Compositional``), and ``ps.lift.generate`` asks NeqSim, for every point of the axes,
which bottom-hole pressure lifts that stream to that tubing-head pressure. What comes back is an
ordinary ``ps.lift.ProductionTable``: the producer's control names it, the deck writes it, and Flow
reads it like any other.

The model is a plain dead-oil model Flow runs; the component description is only the stream NeqSim
flows up the tubing. Every number is in the model's units (METRIC here: m, barsa, degC, sm3/day).

NeqSim runs as a separate program, the runner ``PETROSCRIPT_NEQSIM`` names (docs/install.md says
how to set one up). Without one the script writes the model's deck without the table and says so.

Flow keywords it exercises: VFPPROD from NeqSim's table, WCONPROD with a THP limit and its table
number, and the well's THP summary.

Run it with ``python examples/lift_neqsim.py``; it writes ``out/lift_neqsim/``.
"""

import os

import polars as pl

import petroscript as ps

# 1. The model the table is for: an ordinary model that checks clean. Its units are the table's.
m = ps.Model(
    name="LIFT_NEQSIM",
    units="metric",
    start="2025-01-01",
    grid=ps.grid.Cartesian(shape=(10, 10, 3), dx=300.0, dy=300.0, dz=20.0, top=2000.0),
    rock=ps.Rock(porosity=0.2, permeability=300.0, vertical_ratio=0.1,
                 compressibility=5.0e-5, reference_pressure=230.0),
    fluid=ps.fluid.OilWater(
        oil=ps.pvt.DeadOil(pressure=[50.0, 230.0, 400.0], oil_fvf=[1.12, 1.10, 1.08],
                           oil_viscosity=[1.8, 2.0, 2.3]),
        water=ps.pvt.Water(reference_pressure=230.0, fvf=1.01, compressibility=4.5e-5,
                           viscosity=0.4, viscosibility=0.0),
        oil_surface_density=850.0,
        water_surface_density=1020.0,
    ),
    relperm=[ps.relperm.OilWater(sw=[0.2, 0.4, 0.6, 0.8], krw=[0.0, 0.05, 0.2, 0.5],
                                 krow=[0.9, 0.4, 0.1, 0.0], pcow=[0.0, 0.0, 0.0, 0.0])],
    initial=ps.initial.Equilibrium(datum_depth=2000.0, datum_pressure=230.0,
                                   oil_water_contact=2100.0),
    # the table's pressures refer to this depth: the well's own BHP reference depth
    wells=[ps.Well(name="PROD", i=5, j=5, bhp_reference_depth=2000.0)],
    schedule=ps.Schedule(
        end="2028-01-01",
        report="monthly",
        events=[ps.Event("2025-01-01",
                         ps.Perforate(well="PROD", k=(1, 3), diameter=0.2),
                         ps.Produce(well="PROD", liquid_rate=400.0, bhp_limit=50.0))],
    ),
    summary=ps.Summary(field=["oil_production_rate", "pressure"],
                       wells={"PROD": ["liquid_production_rate", "tubing_head_pressure",
                                       "bottom_hole_pressure"]}),
)
assert not [issue for issue in m.check() if issue.severity == "error"]

# 2. The stream NeqSim flows: the oil by its components, with the composition it is produced at.
#    Water is a component too (H2O): a water-cut axis above 0 needs it in the composition.
oil = ps.fluid.Compositional(
    equation_of_state="peng_robinson",
    components=ps.pvt.Components(
        name=["C1", "C2", "C3", "NC4", "C7P", "C15P", "H2O"],
        molecular_weight=[16.043, 30.07, 44.097, 58.123, 100.0, 210.0, 18.015],      # g/mol
        critical_temperature=[190.56, 305.32, 369.83, 425.12, 540.0, 700.0, 647.3],  # K
        critical_pressure=[45.99, 48.72, 42.48, 37.96, 27.4, 15.0, 220.89],          # barsa
        critical_volume=[0.0986, 0.1455, 0.2, 0.255, 0.4, 0.8, 0.056],               # m3/kg-mol
        acentric_factor=[0.011, 0.099, 0.152, 0.2, 0.35, 0.65, 0.344],
    ),
    binary_interaction_coefficients={("C1", "C7P"): 0.03},
    compositions=[ps.pvt.Composition(
        name="RES",
        mole_fraction={"C1": 0.4, "C2": 0.05, "C3": 0.05, "NC4": 0.04, "C7P": 0.18,
                       "C15P": 0.18, "H2O": 0.1},
    )],
    surface_temperature=15.0,          # degC: the conditions surface volumes are counted at
    surface_pressure=1.01325,          # barsa
)

# 3. The tubing: stations down the string (the first at the wellhead), one bore per segment.
tubing = ps.lift.Tubing(
    measured_depth=[0.0, 2000.0],      # m along the string
    true_vertical_depth=[0.0, 2000.0], # m, in the model's depth reference
    inner_diameter=[0.124],            # m
    roughness=[1.5e-5],                # m
    wellhead_temperature=40.0,         # degC
    temperature_gradient=0.03,         # degC warmer per m down
)

if not os.environ.get("PETROSCRIPT_NEQSIM"):
    ps.compile(m).write("out/lift_neqsim")
    print("no simulator found for NeqSim (set PETROSCRIPT_NEQSIM to a NeqSim runner); the "
          "model's deck, without the table, is in out/lift_neqsim")
    raise SystemExit

# 4. The table: NeqSim solves every point of the axes, each in the model's units.
table = ps.lift.generate(
    tubing,
    model=m, well="PROD", fluid=oil, composition="RES",
    correlation="beggs_brill",         # or "hagedorn_brown", "gray", "mukherjee_brill"
    number=1,
    rate=[50.0, 200.0, 400.0, 800.0], rate_kind="liquid",       # sm3/day of liquid
    thp=[10.0, 20.0, 40.0, 80.0],                              # barsa
    water_fraction=[0.0, 0.5],                                 # water cut, a fraction
    gas_fraction=[0.0],                                        # gas-oil ratio, sm3/sm3
    folder="out/lift_neqsim/neqsim",   # NeqSim's request, result and log
)
print(table.source)                    # NeqSim's version, the correlation, the composition, ...
print(table.frame().head(4))           # one row per point: thp, water_fraction, ..., rate, bhp

# 5. The control names the table itself; the model holds it at lift_tables/1.
m = m.set({
    "lift_tables/1": table,
    "schedule/2025-01-01/PROD/control": ps.Produce(
        well="PROD", liquid_rate=400.0, bhp_limit=50.0, thp_limit=70.0, lift_table=table),
})
deck = ps.compile(m)
deck.write("out/lift_neqsim")

if ps.simulators():
    run = deck.run(folder="out/lift_neqsim/run")
    liquid = run.wells("liquid_production_rate")
    held = run.wells("tubing_head_pressure").filter(pl.col("value") <= 70.0 + 1e-3)
    print("the THP limit takes over on", held["date"][0].date() if len(held) else "no date yet")
    print(f"liquid rate {liquid['value'][0]:.0f} -> {liquid['value'][-1]:.0f} {liquid['unit'][-1]}")
else:
    print("no simulator found (set PETROSCRIPT_FLOW, or put flow on PATH); the deck is in "
          "out/lift_neqsim")