From 63cb103672f95d5b0fa366bcab6827daa031a7f3 Mon Sep 17 00:00:00 2001
From: Mathias Methlie Nilsen <mathias.m.nilsen97@gmail.com>
Date: Wed, 7 Jan 2026 15:40:22 +0100
Subject: [PATCH 1/9] Add JutulDarcyWrapper

---
 src/subsurface/multphaseflow/jutul_darcy.py | 165 ++++++++++++++++++++
 1 file changed, 165 insertions(+)
 create mode 100644 src/subsurface/multphaseflow/jutul_darcy.py

diff --git a/src/subsurface/multphaseflow/jutul_darcy.py b/src/subsurface/multphaseflow/jutul_darcy.py
new file mode 100644
index 0000000..5065dd5
--- /dev/null
+++ b/src/subsurface/multphaseflow/jutul_darcy.py
@@ -0,0 +1,165 @@
+'''
+Simulator wrapper for the JutulDarcy simulator.
+
+This module provides a wrapper interface for running JutulDarcy simulations
+with support for ensemble-based workflows and flexible output formatting.
+'''
+
+#────────────────────────────────────────────────────
+import pandas as pd
+import numpy as np
+import shutil
+import os
+
+from mako.template import Template
+#────────────────────────────────────────────────────
+
+
+__author__ = 'Mathias Methlie Nilsen'
+__all__ = ['JutulDarcyWrapper']
+
+
+class JutulDarcyWrapper:
+
+    def __init__(self, options):
+
+        # Make makofile an mandatory option
+        if 'makofile' not in options:
+            raise ValueError('Wrapper  requires a makofile option')
+        else:
+            self.makofile = options.get('makofile')
+
+        # Other variables
+        self.reporttype = options.get('reporttype', 'days')
+        self.out_format = options.get('out_format', 'list')
+        self.datatype   = options.get('datatype', ['FOPT', 'FGPT', 'FWPT', 'FWIT'])
+        self.parallel   = options.get('parallel', 1)
+        self.platform   = options.get('platform', 'Python')
+        self.datafile   = None
+
+        # This is for PET to work properly (should be removed in future versions)
+        self.input_dict = options
+        self.true_order = [self.reporttype, options['reportpoint']]
+
+
+    def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True) -> dict:
+        '''
+        Run forward simulation for given input parameters.
+
+        Parameters
+        ----------
+        input : dict
+            Input parameters for the simulation.
+
+        idn : int, optional
+            Ensemble member ID, by default 0.
+
+        delete_folder : bool, optional
+            Whether to delete the simulation folder after running, by default True.
+        '''
+
+        # Include ensemble member id in input dict
+        input['member'] = idn
+        
+        # Make simulation folder
+        folder = f'En_{idn}'
+        os.makedirs(folder)
+
+        # Render makofile
+        self.render_makofile(self.makofile, folder, input)
+
+        # Enter simulation folder and run simulation
+        os.chdir(folder)
+
+        if self.platform == 'Python':
+            # Needs to be imported here for multiprocessing to work
+            from jutuldarcy import simulate_data_file
+            res = simulate_data_file(
+                data_file_name=self.datafile, 
+                convert=True, # Convert to output dictionary
+                units='si',   # Use SI units (Sm3 and so on)
+                info_level=-1 # No terminal output
+            )
+        elif self.platform == 'Julia':
+            from jutuldarcy import convert_to_pydict
+            from juliacall import Main as jl
+            jl.seval("using JutulDarcy, Jutul")
+            case  = jl.setup_case_from_data_file(self.datafile)
+            jlres = jl.simulate_reservoir(case, info_level=-1) 
+            res   = convert_to_pydict(jlres, case, units='si')
+
+            # TODO: Make sure the gradient computation works (this example is hardcoded, for a specific case)
+            if False:
+                # Define objective function for gas rate
+                jl.seval("""
+                function objective_function(model, state, dt, step_i, forces)
+                    oil_rate = JutulDarcy.compute_well_qoi(model, state, forces, :PROD, SurfaceOilRateTarget)
+                    return dt*oil_rate
+                end
+                """)
+                
+                # Compute sensitivities with respect to parameters
+                sensitivities = jl.JutulDarcy.reservoir_sensitivities(
+                    case, 
+                    jlres, 
+                    jl.objective_function,
+                    include_parameters=True
+                )
+                
+                # Access permeability gradient
+                poro_grad = sensitivities[jl.Symbol("porosity")]
+                # Convert to numpy array
+                poro_gradient = np.array(poro_grad)
+                
+        os.chdir('..')
+
+        # Delete simulation folder
+        if delete_folder:
+            shutil.rmtree(folder)
+
+        # Extract requested datatypes
+        output = self.extract_datatypes(res, out_format=self.out_format)
+        return output
+
+
+    def render_makofile(self, makofile: str, folder: str, input: dict):
+        '''
+        Render makofile.mako to makofile.DATA using input
+        '''
+        self.datafile = makofile.replace('.mako', '.DATA')
+        template = Template(filename=makofile)
+        with open(os.path.join(folder, self.datafile), 'w') as f:
+            f.write(template.render(**input))
+
+
+    def extract_datatypes(self, res: dict, out_format='list') -> dict:
+        out = {}
+        for orginal_key in self.datatype:
+            key = orginal_key.upper()
+
+            # Check if key is FIELD data
+            if key in res['FIELD']:
+                out[orginal_key] = res['FIELD'][key]
+            # Check if key is WELLS data (format: "DATA:WELL" or "DATA WELL")
+            elif ':' in key or ' ' in key:
+                data_id, well_id = key.replace(':', ' ').split(' ')
+                out[orginal_key] = res['WELLS'][well_id][data_id]
+            else:
+                raise KeyError(f'Data type {key} not found in simulation results')
+        
+        # Format output
+        if out_format == 'list':
+            # Make into a list of dicts where each dict is a time step (pred_data format)
+            out_list = []
+            for i in range(len(res['DAYS'])):
+                time_step_data = {key: np.array([out[key][i]]) for key in out}
+                out_list.append(time_step_data)
+            return out_list
+        
+        elif out_format == 'dict':
+            out['DAYS'] = res['DAYS']
+            return out
+
+        elif out_format == 'dataframe':
+            df = pd.DataFrame(data=out, index=res['DAYS'])
+            return df
\ No newline at end of file

From 33f648d7d9617ffd41e996c3128cbe87e03c7cc2 Mon Sep 17 00:00:00 2001
From: Mathias Methlie Nilsen <mathias.m.nilsen97@gmail.com>
Date: Thu, 8 Jan 2026 09:16:46 +0100
Subject: [PATCH 2/9] Add documentation for JutulDarcyWrapper

---
 src/subsurface/multphaseflow/jutul_darcy.py | 62 ++++++++++++++++++---
 1 file changed, 54 insertions(+), 8 deletions(-)

diff --git a/src/subsurface/multphaseflow/jutul_darcy.py b/src/subsurface/multphaseflow/jutul_darcy.py
index 5065dd5..3909f69 100644
--- a/src/subsurface/multphaseflow/jutul_darcy.py
+++ b/src/subsurface/multphaseflow/jutul_darcy.py
@@ -12,6 +12,7 @@
 import os
 
 from mako.template import Template
+from typing import Union
 #────────────────────────────────────────────────────
 
 
@@ -22,13 +23,38 @@
 class JutulDarcyWrapper:
 
     def __init__(self, options):
+        '''
+        Wrapper for the JutulDarcy simulator [1].
 
+        Parameters
+        ----------
+        options : dict
+            Configuration options for the wrapper.
+            Keys:
+                - 'makofile' or 'runfile': Path to the makofile.mako or runfile.DATA template.
+                - 'reporttype': Type of report (default: 'days').
+                - 'out_format': Output format ('list', 'dict', 'dataframe'; default: 'list').
+                - 'datatype': List of data types to extract (default: ['FOPT', 'FGPT', 'FWPT', 'FWIT']).
+                - 'parallel': Number of parallel simulations (default: 1).
+                - 'platform': 'Python' or 'Julia' (default: 'Python').
+
+        References
+        ----------
+        [1] Møyner, O. (2025).
+            JutulDarcy.jl – a fully differentiable high-performance reservoir simulator
+            based on automatic differentiation. Computational Geosciences, 29, Article 30.
+            https://doi.org/10.1007/s10596-025-10366-6
+        '''
         # Make makofile an mandatory option
-        if 'makofile' not in options:
-            raise ValueError('Wrapper  requires a makofile option')
-        else:
+        if ('makofile' not in options) and ('runfile' not in options):
+            raise ValueError('Wrapper  requires a makofile (or runfile) option')
+        
+        if 'makofile' in options: 
             self.makofile = options.get('makofile')
 
+        if 'runfile' in options:
+            self.makofile = options.get('runfile').split('.')[0] + '.mako'
+
         # Other variables
         self.reporttype = options.get('reporttype', 'days')
         self.out_format = options.get('out_format', 'list')
@@ -42,20 +68,25 @@ def __init__(self, options):
         self.true_order = [self.reporttype, options['reportpoint']]
 
 
-    def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True) -> dict:
+    def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True) -> Union[dict|list|pd.DataFrame]:
         '''
         Run forward simulation for given input parameters.
 
         Parameters
         ----------
-        input : dict
+        input: dict
             Input parameters for the simulation.
 
-        idn : int, optional
+        idn: int, optional
             Ensemble member ID, by default 0.
 
-        delete_folder : bool, optional
+        delete_folder: bool, optional
             Whether to delete the simulation folder after running, by default True.
+
+        Returns
+        -------
+            output: Union[dict, list, pd.DataFrame]
+                Simulation output in the specified format.
         '''
 
         # Include ensemble member id in input dict
@@ -132,7 +163,22 @@ def render_makofile(self, makofile: str, folder: str, input: dict):
             f.write(template.render(**input))
 
 
-    def extract_datatypes(self, res: dict, out_format='list') -> dict:
+    def extract_datatypes(self, res: dict, out_format='list') -> Union[dict|list|pd.DataFrame]:
+        '''
+        Extract requested datatypes from simulation results.
+
+        Parameters
+        ----------
+        res : dict
+            Simulation results dictionary.
+        out_format : str, optional
+            Output format ('list[dict]', 'dict', 'dataframe'), by default 'list'.
+        
+        Returns
+        -------
+            Union[dict, list, pd.DataFrame]
+                Extracted data in the specified format.
+        '''
         out = {}
         for orginal_key in self.datatype:
             key = orginal_key.upper()

From 331e01e117c97371695f1fcc5169af41fb277b84 Mon Sep 17 00:00:00 2001
From: Mathias Methlie Nilsen <mathias.m.nilsen97@gmail.com>
Date: Mon, 12 Jan 2026 11:16:43 +0100
Subject: [PATCH 3/9] Start to inlude adjoints

---
 src/subsurface/multphaseflow/jutul_darcy.py | 164 ++++++++++++++++----
 1 file changed, 137 insertions(+), 27 deletions(-)

diff --git a/src/subsurface/multphaseflow/jutul_darcy.py b/src/subsurface/multphaseflow/jutul_darcy.py
index 3909f69..2f1e31f 100644
--- a/src/subsurface/multphaseflow/jutul_darcy.py
+++ b/src/subsurface/multphaseflow/jutul_darcy.py
@@ -8,6 +8,7 @@
 #────────────────────────────────────────────────────
 import pandas as pd
 import numpy as np
+import warnings
 import shutil
 import os
 
@@ -20,6 +21,14 @@
 __all__ = ['JutulDarcyWrapper']
 
 
+#────────────────────────────────────────────────────────────────────────────────────
+os.environ['PYTHON_JULIACALL_HANDLE_SIGNALS'] = 'yes'
+os.environ['PYTHON_JULIACALL_THREADS'] = 'auto'
+os.environ['PYTHON_JULIACALL_OPTLEVEL'] = '3'
+warnings.filterwarnings('ignore', message='.*juliacall module already imported.*')
+#────────────────────────────────────────────────────────────────────────────────────
+
+
 class JutulDarcyWrapper:
 
     def __init__(self, options):
@@ -62,11 +71,35 @@ def __init__(self, options):
         self.parallel   = options.get('parallel', 1)
         self.platform   = options.get('platform', 'Python')
         self.datafile   = None
+        self.compute_adjoint = False
 
         # This is for PET to work properly (should be removed in future versions)
         self.input_dict = options
         self.true_order = [self.reporttype, options['reportpoint']]
 
+        # Adjoint information
+        if 'well_adjoint_info' in options:
+            self.compute_adjoint = True
+            self.adjoint_info = {'wells': {}}
+
+            for well_id, well_info in options['well_adjoint_info'].items():
+                well_obj = well_info['objective']
+                well_var = well_info['variable']
+
+                if isinstance(well_obj, str):
+                    well_obj = [well_obj]
+                if isinstance(well_var, str):
+                    well_var = [well_var]
+
+                for obj in well_obj:
+                    if not obj in ['mass', 'liquid', 'water', 'oil', 'gas', 'rate']:
+                        raise ValueError(f'Adjoint objective {obj} not supported')
+                    
+                self.adjoint_info['wells'][well_id] = {
+                    'objective': well_obj,
+                    'variable': well_var
+                }
+                
 
     def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True) -> Union[dict|list|pd.DataFrame]:
         '''
@@ -105,42 +138,61 @@ def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True) -> Unio
         if self.platform == 'Python':
             # Needs to be imported here for multiprocessing to work
             from jutuldarcy import simulate_data_file
-            res = simulate_data_file(
+            pyres = simulate_data_file(
                 data_file_name=self.datafile, 
                 convert=True, # Convert to output dictionary
                 units='si',   # Use SI units (Sm3 and so on)
                 info_level=-1 # No terminal output
             )
         elif self.platform == 'Julia':
-            from jutuldarcy import convert_to_pydict
             from juliacall import Main as jl
+            from jutuldarcy import convert_to_pydict
             jl.seval("using JutulDarcy, Jutul")
+
             case  = jl.setup_case_from_data_file(self.datafile)
-            jlres = jl.simulate_reservoir(case, info_level=-1) 
-            res   = convert_to_pydict(jlres, case, units='si')
+            jlres = jl.simulate_reservoir(case, info_level=-1)
+            pyres = convert_to_pydict(jlres, case, units='si')
 
             # TODO: Make sure the gradient computation works (this example is hardcoded, for a specific case)
-            if False:
-                # Define objective function for gas rate
-                jl.seval("""
-                function objective_function(model, state, dt, step_i, forces)
-                    oil_rate = JutulDarcy.compute_well_qoi(model, state, forces, :PROD, SurfaceOilRateTarget)
-                    return dt*oil_rate
-                end
-                """)
-                
-                # Compute sensitivities with respect to parameters
-                sensitivities = jl.JutulDarcy.reservoir_sensitivities(
-                    case, 
-                    jlres, 
-                    jl.objective_function,
-                    include_parameters=True
-                )
+            if self.compute_adjoint:
+
+                # TODO: There might be a better way of structuring the gradient info (this is very preliminary)!
+                gradients = {}
+                for well_id, well_info in self.adjoint_info['wells'].items():
+                    gradients[well_id] = {}
+                    for obj in well_info['objective']:
+                        for var in well_info['variable']:
+
+                            # Define objective function
+                            obj_func = self.well_volume_objective(
+                                well_id=well_id, 
+                                phase=obj, 
+                                jl_import=jl
+                            )
+
+                            # Compute gradients
+                            obj_grad = jl.JutulDarcy.reservoir_sensitivities(
+                                case, 
+                                jlres, 
+                                obj_func,
+                                include_parameters=True,
+                            )
+                            print(self.symdict_to_pydict(obj_grad.data, jl).keys())
+
+                            # Select relevant gradients
+                            if var == 'poro':
+                                poro_grad = np.array(obj_grad[jl.Symbol("porosity")])
+                                gradients[well_id][f'{obj}_of_poro'] = poro_grad
+                            
+                            perm_grad = np.array(obj_grad[jl.Symbol("permeability")])
+                            if var == 'permx':
+                                gradients[well_id][f'{obj}_of_permx'] = perm_grad[0]
+                            if var == 'permy':
+                                gradients[well_id][f'{obj}_of_permy'] = perm_grad[1]
+                            if var == 'permz':
+                                gradients[well_id][f'{obj}_of_permz'] = perm_grad[2]
                 
-                # Access permeability gradient
-                poro_grad = sensitivities[jl.Symbol("porosity")]
-                # Convert to numpy array
-                poro_gradient = np.array(poro_grad)
+                print(gradients)
                 
         os.chdir('..')
 
@@ -149,8 +201,12 @@ def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True) -> Unio
             shutil.rmtree(folder)
 
         # Extract requested datatypes
-        output = self.extract_datatypes(res, out_format=self.out_format)
-        return output
+        output = self.extract_datatypes(pyres, out_format=self.out_format)
+        
+        if self.compute_adjoint:
+            return output, gradients
+        else:
+            return output
 
 
     def render_makofile(self, makofile: str, folder: str, input: dict):
@@ -208,4 +264,58 @@ def extract_datatypes(self, res: dict, out_format='list') -> Union[dict|list|pd.
 
         elif out_format == 'dataframe':
             df = pd.DataFrame(data=out, index=res['DAYS'])
-            return df
\ No newline at end of file
+            return df
+        
+    
+    def well_volume_objective(self, well_id, phase, jl_import):
+        '''
+        Define a well volume objective function for sensitivity analysis.
+
+        Parameters
+        ----------
+        well_id : str
+            Identifier of the well.
+        phase : str
+            Phase type for the well.
+
+        Returns
+        -------
+            jl.objective_function
+                Julia objective function for the specified well and target.
+        '''
+        #jl_import.seval('using JutulDarcy')
+
+        if phase == 'mass':
+            rate = 'TotalSurfaceMassRate'
+        elif phase == 'liquid':
+            rate = 'SurfaceLiquidRateTarget'
+        elif phase == 'water':
+            rate = 'SurfaceWaterRateTarget'
+        elif phase == 'oil':
+            rate = 'SurfaceOilRateTarget'
+        elif phase == 'gas':
+            rate = 'SurfaceGasRateTarget'
+        elif phase == 'rate':
+            rate = 'TotalRateTarget'
+        else:
+            raise ValueError(f'Unknown phase type: {phase}')
+
+        jl_import.seval(f"""
+        function objective_function(model, state, dt, step_i, forces)
+            rate = JutulDarcy.compute_well_qoi(model, state, forces, Symbol("{well_id}"), {rate})
+            return dt*rate
+        end
+        """)
+
+        return jl_import.objective_function
+    
+
+    def symdict_to_pydict(self, symdict, jl_import):
+        '''Convert a Julia symbolic dictionary to a Python dictionary recursively.'''
+        pydict = {}
+        for key, value in symdict.items():
+            if jl_import.isa(value, jl_import.AbstractDict):
+                pydict[str(key)] = self.symdict_to_pydict(value, jl_import)
+            else:
+                pydict[str(key)] = value
+        return pydict
\ No newline at end of file

From ebbcf3f7b53db4566561a2d64f87605afae81fca Mon Sep 17 00:00:00 2001
From: Mathias Methlie Nilsen <mathias.m.nilsen97@gmail.com>
Date: Tue, 20 Jan 2026 09:04:43 +0100
Subject: [PATCH 4/9] Modify JutulDarcyWrapper

---
 src/subsurface/multphaseflow/jutul_darcy.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/src/subsurface/multphaseflow/jutul_darcy.py b/src/subsurface/multphaseflow/jutul_darcy.py
index 2f1e31f..15c697d 100644
--- a/src/subsurface/multphaseflow/jutul_darcy.py
+++ b/src/subsurface/multphaseflow/jutul_darcy.py
@@ -23,10 +23,10 @@
 
 #────────────────────────────────────────────────────────────────────────────────────
 os.environ['PYTHON_JULIACALL_HANDLE_SIGNALS'] = 'yes'
-os.environ['PYTHON_JULIACALL_THREADS'] = 'auto'
+os.environ['PYTHON_JULIACALL_THREADS'] = '1'
 os.environ['PYTHON_JULIACALL_OPTLEVEL'] = '3'
 warnings.filterwarnings('ignore', message='.*juliacall module already imported.*')
-#────────────────────────────────────────────────────────────────────────────────────
+#────────────────────────────────────────────────────────────────────── ──────────────
 
 
 class JutulDarcyWrapper:
@@ -69,7 +69,9 @@ def __init__(self, options):
         self.out_format = options.get('out_format', 'list')
         self.datatype   = options.get('datatype', ['FOPT', 'FGPT', 'FWPT', 'FWIT'])
         self.parallel   = options.get('parallel', 1)
-        self.platform   = options.get('platform', 'Python')
+        self.platform   = options.get('platform', 'Julia')
+        self.units      = options.get('units', 'si')
+
         self.datafile   = None
         self.compute_adjoint = False
 
@@ -141,7 +143,7 @@ def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True) -> Unio
             pyres = simulate_data_file(
                 data_file_name=self.datafile, 
                 convert=True, # Convert to output dictionary
-                units='si',   # Use SI units (Sm3 and so on)
+                units=self.units, 
                 info_level=-1 # No terminal output
             )
         elif self.platform == 'Julia':
@@ -151,7 +153,8 @@ def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True) -> Unio
 
             case  = jl.setup_case_from_data_file(self.datafile)
             jlres = jl.simulate_reservoir(case, info_level=-1)
-            pyres = convert_to_pydict(jlres, case, units='si')
+            #jlres = jl.simulate_reservoir(case)
+            pyres = convert_to_pydict(jlres, case, units=self.units)
 
             # TODO: Make sure the gradient computation works (this example is hardcoded, for a specific case)
             if self.compute_adjoint:

From fa9bf2ae5ede49997b33b5c2553fecceb6f830b2 Mon Sep 17 00:00:00 2001
From: Mathias Methlie Nilsen <mathias.m.nilsen97@gmail.com>
Date: Mon, 26 Jan 2026 13:30:50 +0100
Subject: [PATCH 5/9] Include adjoints

---
 src/subsurface/multphaseflow/jutul_darcy.py | 471 ++++++++++++++------
 1 file changed, 334 insertions(+), 137 deletions(-)

diff --git a/src/subsurface/multphaseflow/jutul_darcy.py b/src/subsurface/multphaseflow/jutul_darcy.py
index 15c697d..33b0263 100644
--- a/src/subsurface/multphaseflow/jutul_darcy.py
+++ b/src/subsurface/multphaseflow/jutul_darcy.py
@@ -14,6 +14,8 @@
 
 from mako.template import Template
 from typing import Union
+from p_tqdm import p_map
+from tqdm import tqdm
 #────────────────────────────────────────────────────
 
 
@@ -28,6 +30,13 @@
 warnings.filterwarnings('ignore', message='.*juliacall module already imported.*')
 #────────────────────────────────────────────────────────────────────── ──────────────
 
+PBAR_OPTS = {
+    'ncols': 110,
+    'colour': "#285475",
+    'bar_format': '{desc}: {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]',
+    'ascii': '-◼', # Custom bar characters for a sleeker look
+}
+
 
 class JutulDarcyWrapper:
 
@@ -45,12 +54,11 @@ def __init__(self, options):
                 - 'out_format': Output format ('list', 'dict', 'dataframe'; default: 'list').
                 - 'datatype': List of data types to extract (default: ['FOPT', 'FGPT', 'FWPT', 'FWIT']).
                 - 'parallel': Number of parallel simulations (default: 1).
-                - 'platform': 'Python' or 'Julia' (default: 'Python').
 
         References
         ----------
         [1] Møyner, O. (2025).
-            JutulDarcy.jl – a fully differentiable high-performance reservoir simulator
+            JutulDarcy.jl - a fully differentiable high-performance reservoir simulator
             based on automatic differentiation. Computational Geosciences, 29, Article 30.
             https://doi.org/10.1007/s10596-025-10366-6
         '''
@@ -69,41 +77,80 @@ def __init__(self, options):
         self.out_format = options.get('out_format', 'list')
         self.datatype   = options.get('datatype', ['FOPT', 'FGPT', 'FWPT', 'FWIT'])
         self.parallel   = options.get('parallel', 1)
-        self.platform   = options.get('platform', 'Julia')
-        self.units      = options.get('units', 'si')
+        self.units      = options.get('units', 'metric') # This is not consistently used!
 
-        self.datafile   = None
-        self.compute_adjoint = False
+        self.datafile = None
+        self.compute_adjoints = False
 
         # This is for PET to work properly (should be removed in future versions)
         self.input_dict = options
         self.true_order = [self.reporttype, options['reportpoint']]
+        self.steps = [i+1 for i in range(len(self.true_order[1]))]
+
+        # Adjoint settings
+        #---------------------------------------------------------------------------------------------------------
+        if 'adjoints' in options:
+            self.compute_adjoints = True
+            self.adjoint_info = {}
+    
+            for datatype in options['adjoints']:
 
-        # Adjoint information
-        if 'well_adjoint_info' in options:
-            self.compute_adjoint = True
-            self.adjoint_info = {'wells': {}}
-
-            for well_id, well_info in options['well_adjoint_info'].items():
-                well_obj = well_info['objective']
-                well_var = well_info['variable']
-
-                if isinstance(well_obj, str):
-                    well_obj = [well_obj]
-                if isinstance(well_var, str):
-                    well_var = [well_var]
-
-                for obj in well_obj:
-                    if not obj in ['mass', 'liquid', 'water', 'oil', 'gas', 'rate']:
-                        raise ValueError(f'Adjoint objective {obj} not supported')
-                    
-                self.adjoint_info['wells'][well_id] = {
-                    'objective': well_obj,
-                    'variable': well_var
-                }
+                # Check that datatype is supported
+                if not datatype in ['mass', 'liquid', 'water', 'oil', 'gas', 'rate']:
+                    raise ValueError(f'Adjoint objective {datatype} not supported')
                 
+                well_id = options['adjoints'][datatype]['well_id']
+                param = options['adjoints'][datatype]['param']
+                unit  = options['adjoints'][datatype]['unit']
+                steps = options['adjoints'][datatype].get('steps', 'all')
+
+                if steps == 'acc':
+                    steps = None
+                elif steps == 'all':
+                    steps = self.steps
+                elif isinstance(steps, int):
+                    steps = [steps]
+
+                if not isinstance(well_id, list):
+                    well_id = [well_id]
+                if not isinstance(param, list):
+                    param = [param]
+
+                for wid in well_id:
+                    self.adjoint_info[f'{datatype}:{wid}'] = {
+                        'datatype': datatype,
+                        'well_id': wid,
+                        'param': param,
+                        'unit': unit,
+                        'steps': steps
+                    }
+        #---------------------------------------------------------------------------------------------------------
+
+    def __call__(self, inputs: dict):
+
+        # Delet all existing En_* folders
+        for item in os.listdir('.'):
+            if os.path.isdir(item) and item.startswith('En_'):
+                shutil.rmtree(item)
+        
+        # simulate all inputs in parallel
+        outputs = p_map(
+            self.run_fwd_sim, 
+            [inputs[n] for n in range(len(inputs))], 
+            list(range(len(inputs))), 
+            num_cpus=self.parallel,
+            unit='sim',
+            **PBAR_OPTS
+        )
+
+        if self.compute_adjoints:
+            results, adjoints = zip(*outputs)
+            return results, adjoints
+        else:
+            return outputs
+                     
 
-    def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True) -> Union[dict|list|pd.DataFrame]:
+    def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True):
         '''
         Run forward simulation for given input parameters.
 
@@ -137,66 +184,92 @@ def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True) -> Unio
         # Enter simulation folder and run simulation
         os.chdir(folder)
 
-        if self.platform == 'Python':
-            # Needs to be imported here for multiprocessing to work
-            from jutuldarcy import simulate_data_file
-            pyres = simulate_data_file(
-                data_file_name=self.datafile, 
-                convert=True, # Convert to output dictionary
-                units=self.units, 
-                info_level=-1 # No terminal output
+        from juliacall import Main as jl
+        from jutuldarcy import convert_to_pydict
+        jl.seval("using JutulDarcy, Jutul")
+
+        # Setup case
+        case = jl.setup_case_from_data_file(self.datafile)
+
+        # Get some grid info
+        nx = case.input_data["GRID"]["cartDims"][0]
+        ny = case.input_data["GRID"]["cartDims"][1]  
+        nz = case.input_data["GRID"]["cartDims"][2]
+        grid = (nx, ny, nz)
+        actnum = np.array(case.input_data["GRID"]["ACTNUM"]) # Shape (nx, ny, nz)
+        actnum_vec = actnum.flatten(order='F')  # Fortran order flattening
+
+        # Simulate and get results
+        jlres = jl.simulate_reservoir(case, info_level=-1)
+        pyres = convert_to_pydict(jlres, case, units=self.units)
+
+        if self.compute_adjoints:
+            adjoints = {key: {} for key in self.adjoint_info}
+
+            pbar = tqdm(
+                adjoints.keys(), 
+                desc=f'Solving adjoints for En_{idn}',
+                position=idn+1,
+                leave=False,
+                unit='obj',
+                dynamic_ncols=False,
+                **PBAR_OPTS
             )
-        elif self.platform == 'Julia':
-            from juliacall import Main as jl
-            from jutuldarcy import convert_to_pydict
-            jl.seval("using JutulDarcy, Jutul")
-
-            case  = jl.setup_case_from_data_file(self.datafile)
-            jlres = jl.simulate_reservoir(case, info_level=-1)
-            #jlres = jl.simulate_reservoir(case)
-            pyres = convert_to_pydict(jlres, case, units=self.units)
-
-            # TODO: Make sure the gradient computation works (this example is hardcoded, for a specific case)
-            if self.compute_adjoint:
-
-                # TODO: There might be a better way of structuring the gradient info (this is very preliminary)!
-                gradients = {}
-                for well_id, well_info in self.adjoint_info['wells'].items():
-                    gradients[well_id] = {}
-                    for obj in well_info['objective']:
-                        for var in well_info['variable']:
-
-                            # Define objective function
-                            obj_func = self.well_volume_objective(
-                                well_id=well_id, 
-                                phase=obj, 
-                                jl_import=jl
-                            )
-
-                            # Compute gradients
-                            obj_grad = jl.JutulDarcy.reservoir_sensitivities(
-                                case, 
-                                jlres, 
-                                obj_func,
-                                include_parameters=True,
-                            )
-                            print(self.symdict_to_pydict(obj_grad.data, jl).keys())
-
-                            # Select relevant gradients
-                            if var == 'poro':
-                                poro_grad = np.array(obj_grad[jl.Symbol("porosity")])
-                                gradients[well_id][f'{obj}_of_poro'] = poro_grad
-                            
-                            perm_grad = np.array(obj_grad[jl.Symbol("permeability")])
-                            if var == 'permx':
-                                gradients[well_id][f'{obj}_of_permx'] = perm_grad[0]
-                            if var == 'permy':
-                                gradients[well_id][f'{obj}_of_permy'] = perm_grad[1]
-                            if var == 'permz':
-                                gradients[well_id][f'{obj}_of_permz'] = perm_grad[2]
-                
-                print(gradients)
+            
+            for key in adjoints:
+                #datatype, well_id, param, unit = self.adjoint_info[key]
+                info = self.adjoint_info[key]
+
+                if info['unit'] == 'rate':
+                    rate = True
+                elif info['unit'] == 'volume':
+                    rate = False
+                else:
+                    raise ValueError(f'Unknown unit type: {info["unit"]}')
+
+                func = get_well_objective(
+                    info['well_id'], 
+                    info['datatype'], 
+                    info['steps'], 
+                    rate=rate, 
+                    jl_import=jl
+                )
+
+                # Define objective function
+                func = func if isinstance(func, list) else [func]
+                grad = []
+                for f in func:
+                    # Compute adjoint gradient
+                    g = jl.JutulDarcy.reservoir_sensitivities(
+                        case, 
+                        jlres, 
+                        f,
+                        include_parameters=True,
+                    )
+                    grad.append(g)
                 
+                # POROSITY
+                if 'poro' in info['param']:
+                    poro_grads = []
+                    for g in grad:
+                        poro_grad = np.array(g[jl.Symbol("porosity")])
+                        poro_grads.append(_expand_to_active_grid(poro_grad, actnum_vec, fill_value=0))
+                    adjoints[key]['poro'] = np.squeeze(np.array(poro_grads))
+
+                # PERMEABILITY
+                m2_per_mD = 9.869233e-16
+                for idx, perm_key in enumerate(['permx', 'permy', 'permz']):
+                    if perm_key in info['param']:
+                        perm_grads = []
+                        for g in grad:
+                            perm_grad = np.array(g[jl.Symbol("permeability")])[idx] * m2_per_mD
+                            perm_grads.append(_expand_to_active_grid(perm_grad, actnum_vec, fill_value=0))
+                        adjoints[key][perm_key] = np.squeeze(np.array(perm_grads))
+
+                # Update progress bar
+                pbar.update(1)
+            pbar.close()
+          
         os.chdir('..')
 
         # Delete simulation folder
@@ -204,14 +277,13 @@ def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True) -> Unio
             shutil.rmtree(folder)
 
         # Extract requested datatypes
-        output = self.extract_datatypes(pyres, out_format=self.out_format)
+        output = self.extract_datatypes(pyres, jlcase=case, out_format=self.out_format)
         
-        if self.compute_adjoint:
-            return output, gradients
+        if self.compute_adjoints:
+            return output, adjoints
         else:
             return output
 
-
     def render_makofile(self, makofile: str, folder: str, input: dict):
         '''
         Render makofile.mako to makofile.DATA using input
@@ -222,7 +294,7 @@ def render_makofile(self, makofile: str, folder: str, input: dict):
             f.write(template.render(**input))
 
 
-    def extract_datatypes(self, res: dict, out_format='list') -> Union[dict|list|pd.DataFrame]:
+    def extract_datatypes(self, res: dict, jlcase =None, out_format='list') -> Union[dict|list|pd.DataFrame]:
         '''
         Extract requested datatypes from simulation results.
 
@@ -249,6 +321,13 @@ def extract_datatypes(self, res: dict, out_format='list') -> Union[dict|list|pd.
             elif ':' in key or ' ' in key:
                 data_id, well_id = key.replace(':', ' ').split(' ')
                 out[orginal_key] = res['WELLS'][well_id][data_id]
+
+            elif key in [str(k) for k in jlcase.input_data["GRID"].keys()]:
+                value = jlcase.input_data["GRID"][f"{key}"]
+                try:
+                    out[orginal_key] = np.array(value)
+                except:
+                    out[orginal_key] = value
             else:
                 raise KeyError(f'Data type {key} not found in simulation results')
         
@@ -269,56 +348,174 @@ def extract_datatypes(self, res: dict, out_format='list') -> Union[dict|list|pd.
             df = pd.DataFrame(data=out, index=res['DAYS'])
             return df
         
-    
-    def well_volume_objective(self, well_id, phase, jl_import):
-        '''
-        Define a well volume objective function for sensitivity analysis.
-
-        Parameters
-        ----------
-        well_id : str
-            Identifier of the well.
-        phase : str
-            Phase type for the well.
 
-        Returns
-        -------
-            jl.objective_function
-                Julia objective function for the specified well and target.
-        '''
-        #jl_import.seval('using JutulDarcy')
-
-        if phase == 'mass':
-            rate = 'TotalSurfaceMassRate'
-        elif phase == 'liquid':
-            rate = 'SurfaceLiquidRateTarget'
-        elif phase == 'water':
-            rate = 'SurfaceWaterRateTarget'
-        elif phase == 'oil':
-            rate = 'SurfaceOilRateTarget'
-        elif phase == 'gas':
-            rate = 'SurfaceGasRateTarget'
-        elif phase == 'rate':
-            rate = 'TotalRateTarget'
+    
+def _symdict_to_pydict(symdict, jl_import):
+    '''Convert a Julia symbolic dictionary to a Python dictionary recursively.'''
+    pydict = {}
+    for key, value in symdict.items():
+        if jl_import.isa(value, jl_import.AbstractDict):
+            pydict[str(key)] = _symdict_to_pydict(value, jl_import)
         else:
-            raise ValueError(f'Unknown phase type: {phase}')
-
+            pydict[str(key)] = value
+    return pydict
+    
+def _expand_to_active_grid(param, active, fill_value=np.nan):
+
+    if len(param) == active.sum():
+        val = []
+        i = 0
+        for cell in active:
+            if cell == 1:
+                val.append(param[i])
+                i += 1
+            else:
+                val.append(fill_value)
+    elif len(param) == len(active):
+        val = param
+    else:
+        raise ValueError('Parameter length does not match number of active cells')
+    
+    return np.array(val)
+
+
+def get_well_objective(well_id, rate_id, step_id, rate=True, jl_import=None):
+    '''
+    Create a Julia objective function for well-based adjoint sensitivity analysis.
+
+    This function generates JutulDarcy objective functions that compute well quantities
+    of interest (QOI) for specific phases. The objective can target all timesteps,
+    specific timesteps, or a single timestep, and can return either instantaneous
+    rates or cumulative volumes.
+
+    Parameters
+    ----------
+    well_id : str
+        Identifier of the well for which to compute the objective.
+    rate_id : str
+        Phase type identifier. Supported values:
+        - 'mass': Total surface mass rate
+        - 'liquid': Surface liquid rate
+        - 'water': Surface water rate
+        - 'oil': Surface oil rate
+        - 'gas': Surface gas rate
+        - 'rate': Total volumetric rate
+    step_id : int, list, np.ndarray, or None
+        Timestep specification:
+        - None: Compute objective for all timesteps (cumulative)
+        - int: Compute objective for a single specific timestep
+        - list/array: Compute objectives for multiple specific timesteps
+    rate : bool, optional
+        If True (default), returns instantaneous rate at timestep(s).
+        If False, multiplies rate by dt for cumulative volume contribution.
+    jl_import : module, optional
+        Julia Main module from juliacall. If None, will import automatically.
+
+    Returns
+    -------
+    function or list of functions
+        - Single Julia objective function if step_id is None or int
+        - List of Julia objective functions if step_id is a list/array
+
+    Raises
+    ------
+    ValueError
+        If rate_id is not one of the supported phase types.
+
+    Examples
+    --------
+    >>> # Objective for total oil production across all timesteps
+    >>> obj = get_well_objective('PROD1', 'oil', None, rate=False)
+    
+    >>> # Objective for water rate at timestep 10
+    >>> obj = get_well_objective('INJ1', 'water', 10, rate=True)
+    
+    >>> # Objectives for gas rate at multiple timesteps
+    >>> objs = get_well_objective('PROD2', 'gas', [5, 10, 15], rate=True)
+    '''
+
+    if jl_import is None:
+        from juliacall import Main as jl_import
+        jl_import.seval('using JutulDarcy')
+
+    rate_id_map = {
+        'mass': 'TotalSurfaceMassRate',
+        'liquid': 'SurfaceLiquidRateTarget',
+        'water': 'SurfaceWaterRateTarget',
+        'oil': 'SurfaceOilRateTarget',
+        'gas': 'SurfaceGasRateTarget',
+        'rate': 'TotalRateTarget'
+    }
+    if rate_id not in rate_id_map:
+        raise ValueError(f'Unknown rate type: {rate_id}')
+    rate_id = rate_id_map[rate_id]
+
+    if rate:
+        dt = ''
+    else:
+        dt = 'dt*'
+
+    # Case 1: Sum of all timesteps
+    #-----------------------------------------------------------------------------
+    if step_id is None:
         jl_import.seval(f"""
         function objective_function(model, state, dt, step_i, forces)
-            rate = JutulDarcy.compute_well_qoi(model, state, forces, Symbol("{well_id}"), {rate})
-            return dt*rate
+            rate = JutulDarcy.compute_well_qoi(
+                model, 
+                state, 
+                forces, 
+                Symbol("{well_id}"), 
+                {rate_id}
+            )
+            return {dt}rate
         end
         """)
-
         return jl_import.objective_function
+    #-----------------------------------------------------------------------------
     
-
-    def symdict_to_pydict(self, symdict, jl_import):
-        '''Convert a Julia symbolic dictionary to a Python dictionary recursively.'''
-        pydict = {}
-        for key, value in symdict.items():
-            if jl_import.isa(value, jl_import.AbstractDict):
-                pydict[str(key)] = self.symdict_to_pydict(value, jl_import)
-            else:
-                pydict[str(key)] = value
-        return pydict
\ No newline at end of file
+    # Case 2: Multiple timesteps
+    #-----------------------------------------------------------------------------
+    elif isinstance(step_id, (list, np.ndarray)):
+        objective_steps = []
+        for sid in step_id:
+            jl_import.seval(f"""
+            function objective_function_{sid}(model, state, dt, step_i, forces)
+                if step_i != {sid}
+                    return 0.0
+                else
+                    rate = JutulDarcy.compute_well_qoi(
+                        model, 
+                        state, 
+                        forces, 
+                        Symbol("{well_id}"), 
+                        {rate_id}
+                    )
+                    return {dt}rate
+                end
+            end
+            """)
+            objective_steps.append(jl_import.seval(f'objective_function_{sid}'))
+        return objective_steps
+    #-----------------------------------------------------------------------------
+
+    # Case 3: Single timestep
+    #-----------------------------------------------------------------------------
+    else:
+        jl_import.seval(f"""
+        function objective_function(model, state, dt, step_i, forces)
+            if step_i != {step_id}
+                return 0.0
+            else
+                rate = JutulDarcy.compute_well_qoi(
+                    model, 
+                    state, 
+                    forces, 
+                    Symbol("{well_id}"), 
+                    {rate_id}
+                )
+                return {dt}rate
+            end
+        end
+        """)
+        return jl_import.objective_function
+    #-----------------------------------------------------------------------------

From cb0a1253633c51585b80e7a9e2b8cb1789b5b5de Mon Sep 17 00:00:00 2001
From: Mathias Methlie Nilsen <mathias.m.nilsen97@gmail.com>
Date: Tue, 27 Jan 2026 13:56:29 +0100
Subject: [PATCH 6/9] Improve adjoint input/output

---
 src/subsurface/multphaseflow/jutul_darcy.py | 172 +++++++++++---------
 1 file changed, 98 insertions(+), 74 deletions(-)

diff --git a/src/subsurface/multphaseflow/jutul_darcy.py b/src/subsurface/multphaseflow/jutul_darcy.py
index 33b0263..6e6b420 100644
--- a/src/subsurface/multphaseflow/jutul_darcy.py
+++ b/src/subsurface/multphaseflow/jutul_darcy.py
@@ -91,38 +91,57 @@ def __init__(self, options):
         #---------------------------------------------------------------------------------------------------------
         if 'adjoints' in options:
             self.compute_adjoints = True
+
             self.adjoint_info = {}
-    
             for datatype in options['adjoints']:
-
-                # Check that datatype is supported
-                if not datatype in ['mass', 'liquid', 'water', 'oil', 'gas', 'rate']:
-                    raise ValueError(f'Adjoint objective {datatype} not supported')
                 
-                well_id = options['adjoints'][datatype]['well_id']
-                param = options['adjoints'][datatype]['param']
-                unit  = options['adjoints'][datatype]['unit']
-                steps = options['adjoints'][datatype].get('steps', 'all')
-
+                # Determine if rate or volume and phase. 
+                if datatype in ['WOPT', 'WGPT', 'WWPT', 'WLPT']:
+                    rate  = False
+                    phase = {
+                        'WOPT': 'oil',
+                        'WGPT': 'gas',
+                        'WWPT': 'water',
+                        'WLPT': 'liquid',
+                    }[datatype]
+                
+                elif datatype in ['WOPR', 'WGPR', 'WWPR', 'WLPR']:
+                    rate  = True
+                    phase = {
+                        'WOPR': 'oil',
+                        'WGPR': 'gas',
+                        'WWPR': 'water',
+                        'WLPR': 'liquid',
+                    }[datatype]
+
+
+                # Determine steps
+                steps = options['adjoints'][datatype].get('steps', 'acc')
                 if steps == 'acc':
-                    steps = None
+                    steps = [self.steps[-1]]
+                    accumulative = True
                 elif steps == 'all':
                     steps = self.steps
+                    accumulative = False
                 elif isinstance(steps, int):
+                    accumulative = False
                     steps = [steps]
+                
+                well_ids = options['adjoints'][datatype]['well_id']
+                parameters = options['adjoints'][datatype]['parameters']
 
-                if not isinstance(well_id, list):
-                    well_id = [well_id]
-                if not isinstance(param, list):
-                    param = [param]
+                # Ensure well_ids and parameters are lists
+                well_ids = well_ids if isinstance(well_ids, (list, tuple)) else [well_ids]
+                parameters = parameters if isinstance(parameters, (list, tuple)) else [parameters]
 
-                for wid in well_id:
+                for wid in well_ids:
                     self.adjoint_info[f'{datatype}:{wid}'] = {
-                        'datatype': datatype,
+                        'rate': rate,
+                        'phase': phase,
                         'well_id': wid,
-                        'param': param,
-                        'unit': unit,
-                        'steps': steps
+                        'parameters': parameters,
+                        'steps': steps,
+                        'accumulative': accumulative,
                     }
         #---------------------------------------------------------------------------------------------------------
 
@@ -170,6 +189,9 @@ def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True):
             output: Union[dict, list, pd.DataFrame]
                 Simulation output in the specified format.
         '''
+        from juliacall import Main as jl
+        from jutuldarcy import convert_to_pydict
+        jl.seval("using JutulDarcy, Jutul")
 
         # Include ensemble member id in input dict
         input['member'] = idn
@@ -184,10 +206,6 @@ def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True):
         # Enter simulation folder and run simulation
         os.chdir(folder)
 
-        from juliacall import Main as jl
-        from jutuldarcy import convert_to_pydict
-        jl.seval("using JutulDarcy, Jutul")
-
         # Setup case
         case = jl.setup_case_from_data_file(self.datafile)
 
@@ -204,8 +222,17 @@ def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True):
         pyres = convert_to_pydict(jlres, case, units=self.units)
 
         if self.compute_adjoints:
-            adjoints = {key: {} for key in self.adjoint_info}
 
+            # Initialize adjoint dataframe
+            colnames = []
+            for key in self.adjoint_info:
+                for param in self.adjoint_info[key]['parameters']:
+                    colnames.append((key, param))
+
+            adjoints = pd.DataFrame(columns=pd.MultiIndex.from_tuples(colnames), index=self.true_order[1])
+            adjoints.index.name = self.true_order[0]
+
+            # Initialize progress bar
             pbar = tqdm(
                 adjoints.keys(), 
                 desc=f'Solving adjoints for En_{idn}',
@@ -216,60 +243,63 @@ def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True):
                 **PBAR_OPTS
             )
             
-            for key in adjoints:
-                #datatype, well_id, param, unit = self.adjoint_info[key]
-                info = self.adjoint_info[key]
-
-                if info['unit'] == 'rate':
-                    rate = True
-                elif info['unit'] == 'volume':
-                    rate = False
-                else:
-                    raise ValueError(f'Unknown unit type: {info["unit"]}')
-
-                func = get_well_objective(
-                    info['well_id'], 
-                    info['datatype'], 
-                    info['steps'], 
-                    rate=rate, 
+            # Loop over adjoint objectives
+            for col in adjoints.columns.levels[0]:
+                info = self.adjoint_info[col]
+
+                funcs = get_well_objective(
+                    well_id=info['well_id'],
+                    rate_id=info['phase'],
+                    step_id=info['steps'],
+                    rate=info['rate'],
+                    accumulative=info['accumulative'],
                     jl_import=jl
                 )
 
                 # Define objective function
-                func = func if isinstance(func, list) else [func]
-                grad = []
-                for f in func:
+                funcs = funcs if isinstance(funcs, list) else [funcs]
+                grads = []
+                for func in funcs:
                     # Compute adjoint gradient
-                    g = jl.JutulDarcy.reservoir_sensitivities(
+                    grad = jl.JutulDarcy.reservoir_sensitivities(
                         case, 
                         jlres, 
-                        f,
+                        func,
                         include_parameters=True,
                     )
-                    grad.append(g)
+                    grads.append(grad)
+
+                # Extract and store gradients in adjoint dataframe
+                for g, grad in enumerate(grads):
+                    for param in info['parameters']:
+                        index = self.true_order[1][info['steps'][g]-1]
+                        
+                        if param.lower() == 'poro':
+                            grad_param = np.array(grad[jl.Symbol("porosity")])
+                            grad_param = _expand_to_active_grid(grad_param, actnum_vec, fill_value=0)
+                            adjoints.at[index, (col, param)] = grad_param
+
+                        elif param.lower().startswith('perm'):
+                            grad_param = np.array(grad[jl.Symbol("permeability")])
+
+                            m2_per_mD = 9.869233e-16
+                            if param.lower() == 'permx':
+                                grad_param = grad_param[0] * m2_per_mD
+                            elif param.lower() == 'permy':
+                                grad_param = grad_param[1] * m2_per_mD
+                            elif param.lower() == 'permz':
+                                grad_param = grad_param[2] * m2_per_mD
+                            
+                            grad_param = _expand_to_active_grid(grad_param, actnum_vec, fill_value=0)
+                            adjoints.at[index, (col, param)] = grad_param
+                        else:
+                            raise ValueError(f'Param: {param} not supported for adjoint sensitivity')
+                
                 
-                # POROSITY
-                if 'poro' in info['param']:
-                    poro_grads = []
-                    for g in grad:
-                        poro_grad = np.array(g[jl.Symbol("porosity")])
-                        poro_grads.append(_expand_to_active_grid(poro_grad, actnum_vec, fill_value=0))
-                    adjoints[key]['poro'] = np.squeeze(np.array(poro_grads))
-
-                # PERMEABILITY
-                m2_per_mD = 9.869233e-16
-                for idx, perm_key in enumerate(['permx', 'permy', 'permz']):
-                    if perm_key in info['param']:
-                        perm_grads = []
-                        for g in grad:
-                            perm_grad = np.array(g[jl.Symbol("permeability")])[idx] * m2_per_mD
-                            perm_grads.append(_expand_to_active_grid(perm_grad, actnum_vec, fill_value=0))
-                        adjoints[key][perm_key] = np.squeeze(np.array(perm_grads))
-
                 # Update progress bar
                 pbar.update(1)
             pbar.close()
-          
+
         os.chdir('..')
 
         # Delete simulation folder
@@ -348,7 +378,6 @@ def extract_datatypes(self, res: dict, jlcase =None, out_format='list') -> Union
             df = pd.DataFrame(data=out, index=res['DAYS'])
             return df
         
-
     
 def _symdict_to_pydict(symdict, jl_import):
     '''Convert a Julia symbolic dictionary to a Python dictionary recursively.'''
@@ -379,7 +408,7 @@ def _expand_to_active_grid(param, active, fill_value=np.nan):
     return np.array(val)
 
 
-def get_well_objective(well_id, rate_id, step_id, rate=True, jl_import=None):
+def get_well_objective(well_id, rate_id, step_id, rate=True, accumulative=True, jl_import=None):
     '''
     Create a Julia objective function for well-based adjoint sensitivity analysis.
 
@@ -424,13 +453,8 @@ def get_well_objective(well_id, rate_id, step_id, rate=True, jl_import=None):
 
     Examples
     --------
-    >>> # Objective for total oil production across all timesteps
     >>> obj = get_well_objective('PROD1', 'oil', None, rate=False)
-    
-    >>> # Objective for water rate at timestep 10
     >>> obj = get_well_objective('INJ1', 'water', 10, rate=True)
-    
-    >>> # Objectives for gas rate at multiple timesteps
     >>> objs = get_well_objective('PROD2', 'gas', [5, 10, 15], rate=True)
     '''
 
@@ -457,7 +481,7 @@ def get_well_objective(well_id, rate_id, step_id, rate=True, jl_import=None):
 
     # Case 1: Sum of all timesteps
     #-----------------------------------------------------------------------------
-    if step_id is None:
+    if accumulative:
         jl_import.seval(f"""
         function objective_function(model, state, dt, step_i, forces)
             rate = JutulDarcy.compute_well_qoi(

From 6bbba83ece0034818df6f230d4bae2e2760651f1 Mon Sep 17 00:00:00 2001
From: Mathias Methlie Nilsen <mathias.m.nilsen97@gmail.com>
Date: Tue, 27 Jan 2026 14:31:24 +0100
Subject: [PATCH 7/9] Fix index bug

---
 src/subsurface/multphaseflow/jutul_darcy.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/subsurface/multphaseflow/jutul_darcy.py b/src/subsurface/multphaseflow/jutul_darcy.py
index 6e6b420..605a9df 100644
--- a/src/subsurface/multphaseflow/jutul_darcy.py
+++ b/src/subsurface/multphaseflow/jutul_darcy.py
@@ -85,7 +85,7 @@ def __init__(self, options):
         # This is for PET to work properly (should be removed in future versions)
         self.input_dict = options
         self.true_order = [self.reporttype, options['reportpoint']]
-        self.steps = [i+1 for i in range(len(self.true_order[1]))]
+        self.steps = [i for i in range(len(self.true_order[1]))]
 
         # Adjoint settings
         #---------------------------------------------------------------------------------------------------------
@@ -272,7 +272,7 @@ def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True):
                 # Extract and store gradients in adjoint dataframe
                 for g, grad in enumerate(grads):
                     for param in info['parameters']:
-                        index = self.true_order[1][info['steps'][g]-1]
+                        index = self.true_order[1][info['steps'][g]]
                         
                         if param.lower() == 'poro':
                             grad_param = np.array(grad[jl.Symbol("porosity")])
@@ -504,7 +504,7 @@ def get_well_objective(well_id, rate_id, step_id, rate=True, accumulative=True,
         for sid in step_id:
             jl_import.seval(f"""
             function objective_function_{sid}(model, state, dt, step_i, forces)
-                if step_i != {sid}
+                if step_i != {sid+1}
                     return 0.0
                 else
                     rate = JutulDarcy.compute_well_qoi(
@@ -527,7 +527,7 @@ def get_well_objective(well_id, rate_id, step_id, rate=True, accumulative=True,
     else:
         jl_import.seval(f"""
         function objective_function(model, state, dt, step_i, forces)
-            if step_i != {step_id}
+            if step_i != {step_id+1}
                 return 0.0
             else
                 rate = JutulDarcy.compute_well_qoi(

From 9569b50b59eeb897b6af2ec260835a26c7711bd8 Mon Sep 17 00:00:00 2001
From: Mathias Methlie Nilsen <mathias.m.nilsen97@gmail.com>
Date: Tue, 27 Jan 2026 14:37:29 +0100
Subject: [PATCH 8/9] fix

---
 src/subsurface/multphaseflow/jutul_darcy.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/subsurface/multphaseflow/jutul_darcy.py b/src/subsurface/multphaseflow/jutul_darcy.py
index 605a9df..5940aa3 100644
--- a/src/subsurface/multphaseflow/jutul_darcy.py
+++ b/src/subsurface/multphaseflow/jutul_darcy.py
@@ -28,7 +28,7 @@
 os.environ['PYTHON_JULIACALL_THREADS'] = '1'
 os.environ['PYTHON_JULIACALL_OPTLEVEL'] = '3'
 warnings.filterwarnings('ignore', message='.*juliacall module already imported.*')
-#────────────────────────────────────────────────────────────────────── ──────────────
+#────────────────────────────────────────────────────────────────────────────────────
 
 PBAR_OPTS = {
     'ncols': 110,
@@ -87,7 +87,7 @@ def __init__(self, options):
         self.true_order = [self.reporttype, options['reportpoint']]
         self.steps = [i for i in range(len(self.true_order[1]))]
 
-        # Adjoint settings
+        # Extract adjoint options
         #---------------------------------------------------------------------------------------------------------
         if 'adjoints' in options:
             self.compute_adjoints = True

From 84b782b00a52078085da0c29dd3e54b00cf3e80a Mon Sep 17 00:00:00 2001
From: Mathias Methlie Nilsen <mathias.m.nilsen97@gmail.com>
Date: Wed, 28 Jan 2026 10:48:39 +0100
Subject: [PATCH 9/9] Consistent unit handling

---
 src/subsurface/multphaseflow/jutul_darcy.py | 173 ++++++++++++++------
 1 file changed, 119 insertions(+), 54 deletions(-)

diff --git a/src/subsurface/multphaseflow/jutul_darcy.py b/src/subsurface/multphaseflow/jutul_darcy.py
index 5940aa3..61ac95d 100644
--- a/src/subsurface/multphaseflow/jutul_darcy.py
+++ b/src/subsurface/multphaseflow/jutul_darcy.py
@@ -78,10 +78,24 @@ def __init__(self, options):
         self.datatype   = options.get('datatype', ['FOPT', 'FGPT', 'FWPT', 'FWIT'])
         self.parallel   = options.get('parallel', 1)
         self.units      = options.get('units', 'metric') # This is not consistently used!
-
         self.datafile = None
         self.compute_adjoints = False
 
+        # Process datatypes
+        if isinstance(self.datatype, list) and len(self.datatype) == 1 and isinstance(self.datatype[0], dict):
+            d = []
+            for key in self.datatype[0]:
+                if self.datatype[0][key] is None:
+                    d.append(key)
+                else:
+                    for well in self.datatype[0][key]:
+                        d.append(f'{key}:{well}')
+            self.datatype = d
+        elif isinstance(self.datatype, list) and all(isinstance(dt, str) for dt in self.datatype):
+            self.datatype = self.datatype
+        else:
+            assert self.datatype == list, 'datatype must be a list or dict(list)'
+
         # This is for PET to work properly (should be removed in future versions)
         self.input_dict = options
         self.true_order = [self.reporttype, options['reportpoint']]
@@ -117,6 +131,8 @@ def __init__(self, options):
 
                 # Determine steps
                 steps = options['adjoints'][datatype].get('steps', 'acc')
+                accumulative= False
+
                 if steps == 'acc':
                     steps = [self.steps[-1]]
                     accumulative = True
@@ -164,6 +180,9 @@ def __call__(self, inputs: dict):
 
         if self.compute_adjoints:
             results, adjoints = zip(*outputs)
+            if len(inputs) == 1:
+                results  = results[0]
+                adjoints = adjoints[0]
             return results, adjoints
         else:
             return outputs
@@ -219,8 +238,22 @@ def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True):
 
         # Simulate and get results
         jlres = jl.simulate_reservoir(case, info_level=-1)
-        pyres = convert_to_pydict(jlres, case, units=self.units)
+        pyres = convert_to_pydict(jlres, case, units='metric')
+        pyres = self.results_to_dataframe(pyres, self.datatype, jlcase=case, jl_import=jl)
+
+        # Convert output to requested format
+        if not self.out_format == 'dataframe':
+            if self.out_format == 'dict':
+                output = pyres.to_dict(orient='list')
+            elif self.out_format == 'list':
+                output = []
+                for idx, row in pyres.iterrows():
+                    output.append(row.to_dict())
+        else:
+            output = pyres
 
+        
+        # Compute adjoints
         if self.compute_adjoints:
 
             # Initialize adjoint dataframe
@@ -231,6 +264,7 @@ def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True):
 
             adjoints = pd.DataFrame(columns=pd.MultiIndex.from_tuples(colnames), index=self.true_order[1])
             adjoints.index.name = self.true_order[0]
+            attrs = {}
 
             # Initialize progress bar
             pbar = tqdm(
@@ -278,17 +312,26 @@ def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True):
                             grad_param = np.array(grad[jl.Symbol("porosity")])
                             grad_param = _expand_to_active_grid(grad_param, actnum_vec, fill_value=0)
                             adjoints.at[index, (col, param)] = grad_param
+                            attrs[(col, param)] = {'unit': 'Sm3'}
 
                         elif param.lower().startswith('perm'):
                             grad_param = np.array(grad[jl.Symbol("permeability")])
+                            mD_per_m2 = _convert_from_si(1.0, 'darcy', jl) * 1e3
+                            grad_param  = grad_param/mD_per_m2  # Convert from m2 to mD
+
+                            if info['rate']:
+                                days_per_sec = _convert_from_si(1.0, 'day', jl)
+                                grad_param = grad_param/days_per_sec
+                                attrs[(col, param)] = {'unit': 'Sm3/(day∙mD)'}
+                            else:
+                                attrs[(col, param)] = {'unit': 'Sm3/mD'}
 
-                            m2_per_mD = 9.869233e-16
                             if param.lower() == 'permx':
-                                grad_param = grad_param[0] * m2_per_mD
+                                grad_param = grad_param[0]
                             elif param.lower() == 'permy':
-                                grad_param = grad_param[1] * m2_per_mD
+                                grad_param = grad_param[1]
                             elif param.lower() == 'permz':
-                                grad_param = grad_param[2] * m2_per_mD
+                                grad_param = grad_param[2]
                             
                             grad_param = _expand_to_active_grid(grad_param, actnum_vec, fill_value=0)
                             adjoints.at[index, (col, param)] = grad_param
@@ -299,15 +342,14 @@ def run_fwd_sim(self, input: dict, idn: int=0, delete_folder: bool=True):
                 # Update progress bar
                 pbar.update(1)
             pbar.close()
+            adjoints.attrs = attrs
 
         os.chdir('..')
+        
 
         # Delete simulation folder
         if delete_folder:
             shutil.rmtree(folder)
-
-        # Extract requested datatypes
-        output = self.extract_datatypes(pyres, jlcase=case, out_format=self.out_format)
         
         if self.compute_adjoints:
             return output, adjoints
@@ -324,60 +366,47 @@ def render_makofile(self, makofile: str, folder: str, input: dict):
             f.write(template.render(**input))
 
 
-    def extract_datatypes(self, res: dict, jlcase =None, out_format='list') -> Union[dict|list|pd.DataFrame]:
-        '''
-        Extract requested datatypes from simulation results.
-
-        Parameters
-        ----------
-        res : dict
-            Simulation results dictionary.
-        out_format : str, optional
-            Output format ('list[dict]', 'dict', 'dataframe'), by default 'list'.
+    def results_to_dataframe(self, res: dict, datatypes: list, jlcase=None, jl_import=None) -> pd.DataFrame:
         
-        Returns
-        -------
-            Union[dict, list, pd.DataFrame]
-                Extracted data in the specified format.
-        '''
-        out = {}
-        for orginal_key in self.datatype:
-            key = orginal_key.upper()
+        df = pd.DataFrame(columns=datatypes, index=self.true_order[1])
+        df.index.name = self.true_order[0]
+        attrs = {}
+
+        for key in datatypes:
+            key_upper = key.upper()
 
             # Check if key is FIELD data
-            if key in res['FIELD']:
-                out[orginal_key] = res['FIELD'][key]
+            if key_upper in res['FIELD']:
+                df[key] = res['FIELD'][key_upper]
+                attrs[key_upper] = {'unit': _metric_unit(key_upper)}
+
             # Check if key is WELLS data (format: "DATA:WELL" or "DATA WELL")
-            elif ':' in key or ' ' in key:
-                data_id, well_id = key.replace(':', ' ').split(' ')
-                out[orginal_key] = res['WELLS'][well_id][data_id]
+            elif ':' in key_upper or ' ' in key_upper:
+                data_id, well_id = key_upper.replace(':', ' ').split(' ')
+                df[key] = res['WELLS'][well_id][data_id]
+                attrs[key_upper] = {'unit': _metric_unit(data_id)}
+
+            elif key_upper in [str(k) for k in jlcase.input_data["GRID"].keys()] and (jlcase is not None):
+                value = jlcase.input_data["GRID"][f"{key_upper}"]
+                
+                if key_upper.startswith('PERM'):
+                    value = _convert_from_si(value, 'darcy', jl_import)
+                    value = np.array(value) * 1e3 # Darcy to mD
+                    attrs[key_upper] = {'unit': 'mD'}
+                else:
+                    attrs[key_upper] = {'unit': _metric_unit(key_upper)}
 
-            elif key in [str(k) for k in jlcase.input_data["GRID"].keys()]:
-                value = jlcase.input_data["GRID"][f"{key}"]
                 try:
-                    out[orginal_key] = np.array(value)
+                    df.at[df.index[0], key] = np.array(value)
                 except:
-                    out[orginal_key] = value
+                    df.at[df.index[0], key] = value
+                
             else:
                 raise KeyError(f'Data type {key} not found in simulation results')
-        
-        # Format output
-        if out_format == 'list':
-            # Make into a list of dicts where each dict is a time step (pred_data format)
-            out_list = []
-            for i in range(len(res['DAYS'])):
-                time_step_data = {key: np.array([out[key][i]]) for key in out}
-                out_list.append(time_step_data)
-            return out_list
-        
-        elif out_format == 'dict':
-            out['DAYS'] = res['DAYS']
-            return out
-
-        elif out_format == 'dataframe':
-            df = pd.DataFrame(data=out, index=res['DAYS'])
-            return df
-        
+            
+        df.attrs = attrs
+        return df
+    
     
 def _symdict_to_pydict(symdict, jl_import):
     '''Convert a Julia symbolic dictionary to a Python dictionary recursively.'''
@@ -408,6 +437,42 @@ def _expand_to_active_grid(param, active, fill_value=np.nan):
     return np.array(val)
 
 
+def _convert_from_si(value, unit, jl_import):
+    '''Convert value from SI units to specified unit using JutulDarcy conversion.'''
+    return jl_import.Jutul.convert_from_si(value, jl_import.Symbol(unit))
+
+def _metric_unit(key: str) -> str:
+    '''Return metric unit for given key.'''
+    unit_map = {
+        'PORO': '',
+        'PERMX': 'mD',
+        'PERMY': 'mD',
+        'PERMZ': 'mD',
+        #---------------------
+        'FOPT': 'Sm3',
+        'FGPT': 'Sm3',
+        'FWPT': 'Sm3',
+        'FWLT': 'Sm3',
+        'FWIT': 'Sm3',
+        #---------------------
+        'FOPR': 'Sm3/day',
+        'FGPR': 'Sm3/day',
+        'FWPR': 'Sm3/day',
+        'FLPR': 'Sm3/day',
+        'FWIR': 'Sm3/day',
+        #---------------------
+        'WOPR': 'Sm3/day',
+        'WGPR': 'Sm3/day',
+        'WWPR': 'Sm3/day',
+        'WLPR': 'Sm3/day',
+        'WWIR': 'Sm3/day',
+    }
+    if key.upper() in unit_map:
+        return unit_map[key.upper()]
+    else:
+        return 'Unknown'
+
+
 def get_well_objective(well_id, rate_id, step_id, rate=True, accumulative=True, jl_import=None):
     '''
     Create a Julia objective function for well-based adjoint sensitivity analysis.