Flow solver is now almost fully refactored and optimised

.coveragerc

@@ -4,6 +4,7 @@ omit =
     src/pybella/data_assimilation/*
     src/pybella/inputs/*
     src/pybella/utils/debug_helpers.py
+    src/pybella/utils/slices.py
     */tests/*
     */test_*
     setup.py

.gitignore

@@ -5,8 +5,12 @@
 *.ipynb_checkpoints*
 *.egg-info*
 *.coverage*
-profile.html
-profile.json
+profile*.html
+profile*.json
+profile*.out
+prof/*
+profiling/*
+chat*.json
 
 # Ignore everything under outputs/
 outputs/**
@@ -28,6 +32,7 @@ outputs/**
 *copy.py*
 package*json
 *.log
+*.swp
 
 # Sphinx html build
 **_build*

docs/source/apis/src.flow_solver.physics.low_mach.mpv.rst

@@ -1,7 +1,7 @@
-src.flow_solver.physics.low\_mach.mpv
+src.flow_solver.physics.low\_mach.npf
 ================================
 
-.. automodule:: src.flow_solver.physics.low_mach.mpv
+.. automodule:: src.flow_solver.physics.low_mach.npf
 
 
 

docs/source/apis/src.flow_solver.physics.low_mach.rst

@@ -28,6 +28,6 @@ src.flow_solver.physics.low\_mach
    :recursive:
 
    src.flow_solver.physics.low_mach.laplacian
-   src.flow_solver.physics.low_mach.mpv
+   src.flow_solver.physics.low_mach.npf
    src.flow_solver.physics.low_mach.second_projection
 

docs/source/content/boundary.rst

@@ -4,7 +4,7 @@ Boundary handling
 =================
 Boundary handling for the ghosts cells (:class:`discretization.kgrid.ElemSpaceDiscr`) and nodes (:class:`discretization.kgrid.NodeSpaceDiscr`) are handled in ``boundary.py``.
 
-Cell boundarys are handled by the function :func:`inputs.boundary.set_explicit_boundary_data`. The bondary conditions are given in the initial user data file, :py:attr:`inputs.user_data.UserDataInit.bdry_type`.
+Cell boundarys are handled by the function :func:`inputs.boundary.set_data`. The bondary conditions are given in the initial user data file, :py:attr:`inputs.user_data.UserDataInit.bdry_type`.
 
 For ghost cells in directions without gravity, the ghost cells are padded by the ``np.pad()`` function.
 

src/pybella/__main__.py

@@ -69,7 +69,9 @@ def main():
             debug_writer = io.create_debug_writer(params.debug, writer, mem)
 
             logging.info("For ensemble member = %i..." % cnt)
-            mem = dis_time_update.do(mem, sst.ud, tout, blend, step_writer, debug_writer)
+            mem = dis_time_update.do(
+                mem, sst.ud, tout, blend, step_writer, debug_writer
+            )
 
             if sst.ud.diag:
                 if sst.ud.diag_updt_targets:
@@ -105,8 +107,6 @@ def main():
                 label = "ensemble_mem=%i_%.3f" % (n, tout)
             writer.write_all(mem, str(label) + "_after_full_step")
 
-        # synchronise_variables(mpv, Sol, elem, node, ud, th)
-        # sst.t = tout
         tout_old = np.copy(tout)
         logging.info("tout = %.3f" % tout)
 

src/pybella/data_assimilation/analysis.py

@@ -1,18 +1,17 @@
 import logging
-
 import numpy as np
 
+from ..utils import sim_params as params
+from ..flow_solver.utils.boundary import cell_boundary as bdry_c
+from ..flow_solver.utils.boundary import node_boundary as bdry_n
+
 from . import (
     etpf as da_etpf,
     post_processing as da_post_processing,
     letkf as da_letkf,
     utils as da_utils,
 )
 
-from ..flow_solver.utils import boundary as bdry
-
-from ..utils import sim_params as params
-
 
 def do_for_window(tout, outer_step, results, sst, writer):
     # mp = sst.model_params
@@ -30,9 +29,9 @@ def do_for_window(tout, outer_step, results, sst, writer):
         # Update ensemble with forecast
         ######################################################
         for mem in results:
-            elem, node, sol, _, mpv, th, _ = mem
-            bdry.set_explicit_boundary_data(sol, elem, sst.ud, th, mpv)
-            bdry.set_ghostnodes_p2(mpv.p2_nodes, node, sst.ud)
+            elem, node, sol, _, npf, th, _ = mem
+            bdry_c.set_ghost_cells(mem, sst.ud)
+            bdry_n.set_ghost_nodes(npf.p2_nodes, node, sst.ud)
 
         # ens.set_members(results, tout)
         sst.ensemble_state.set_members(results)
@@ -42,13 +41,13 @@ def do_for_window(tout, outer_step, results, sst, writer):
         ######################################################
         logging.info("Starting output...")
         for mem in sst.ensemble_state:
-            elem, node, sol, _, mpv, th, _ = mem
+            elem, node, sol, _, npf, th, _ = mem
             if params.label_type == "STEP":
                 step = outer_step
                 label = "ensemble_mem=%i_%.3d" % (n, step)
             else:
                 label = "ensemble_mem=%i_%.3f" % (n, tout)
-            writer.write_all(sol, mpv, elem, node, th, str(label) + "_before_da")
+            writer.write_all(sol, npf, elem, node, th, str(label) + "_before_da")
 
         ##################################################
         # LETKF with batch observations
@@ -88,7 +87,7 @@ def do_for_window(tout, outer_step, results, sst, writer):
             )
             results = da_utils.HSprojector_2t3D(results, elem, node, dp.dap, sst.N)
             # if hasattr(dap, 'converter'):
-            # results = dap.converter(results, N, mpv, elem, node, th, ud)
+            # results = dap.converter(results, N, npf, elem, node, th, ud)
 
         ##################################################
         # ETPF
@@ -120,9 +119,9 @@ def do_for_window(tout, outer_step, results, sst, writer):
     # Update ensemble with analysis
     ######################################################
     for mem in results:
-        elem, node, Sol, _, mpv, th, _, _ = mem
-        bdry.set_explicit_boundary_data(Sol, elem, sst.ud, th, mpv)
-        p2_nodes = mpv.p2_nodes
-        bdry.set_ghostnodes_p2(p2_nodes, node, sst.ud)
+        elem, node, sol, npf, th, _, _ = mem
+        bdry_c.set_ghost_cells(mem, sst.ud)
+        p2_nodes = npf.p2_nodes
+        bdry_n.set_ghost_nodes(p2_nodes, node, sst.ud)
 
     sst.ensemble_state.set_members(results)

src/pybella/data_assimilation/params.py

@@ -4,7 +4,8 @@
 import numpy as np
 import scipy as sp
 
-from ..flow_solver.utils import boundary as bdry
+from ..flow_solver.utils.boundary import cell_boundary as bdry_c
+from ..flow_solver.utils.boundary import node_boundary as bdry_n
 
 
 class init(object):
@@ -23,15 +24,6 @@ def __init__(self, N, da_type="rloc"):
         # which attributes to inflate in ensemble inflation?
         self.attributes = ["rho", "rhou", "rhov"]
 
-        # self.obs_path = './output_travelling_vortex/output_travelling_vortex_ensemble=1_32_32_6.0_truthgen.h5'
-        # self.obs_path = './output_rising_bubble/output_rising_bubble_ensemble=1_100_50_10.0_psinc_ref.h5'
-        # self.obs_path = './output_rising_bubble/output_rising_bubble_ensemble=1_100_50_10.0_truthgen_freezelt5.h5'
-        # self.obs_path = './output_rising_bubble/output_rising_bubble_ensemble=1_100_50_10.0_comp_ref.h5'
-        # self.obs_path = './output_rising_bubble/output_rising_bubble_ensemble=1_160_80_1.0_truth_CFLfixed_ib-0.h5'
-
-        # self.obs_path = './output_swe_vortex/output_swe_vortex_ensemble=1_64_1_64_3.0_comp_1.0_pps_tra_truth.h5'
-        # self.obs_path = './output_swe_vortex/output_swe_vortex_ensemble=1_64_1_64_3.0_neg_comp_1.0_pp_tra_truth_ip.h5'
-        # self.obs_path = './output_travelling_vortex/output_travelling_vortex_ensemble=1_64_64_3.0_comp_1.0_pp_tra_truth_ip.h5'
         self.obs_path = "./output_travelling_vortex/output_travelling_vortex_ensemble=1_64_64_3.0_obs.h5"
 
         # forward operator (projector from state space to observation space)
@@ -129,7 +121,7 @@ def gen_obs_noise(self):
             self.obs_noise_seeds = [np.random.randint(10000)]
 
     @staticmethod
-    def converter(results, N, mpv, elem, node, th, ud):
+    def converter(results, N, npf, elem, node, th, ud):
         """
         Do this after data assimilation for HS balanced vortex.
 
@@ -138,7 +130,7 @@ def converter(results, N, mpv, elem, node, th, ud):
 
         g = ud.g0
         for n in range(N):
-            bdry.set_explicit_boundary_data(results[n][0], elem, ud, th, mpv)
+            bdry_c.set_ghost_cells(results[n][0], elem, ud, th, npf)
             results[n][0].rhoY[...] = (g / 2.0 * results[n][0].rho ** 2) ** th.gamminv
 
             igy = elem.igy
@@ -148,12 +140,12 @@ def converter(results, N, mpv, elem, node, th, ud):
 
             pn = sp.signal.convolve(results[n][0].rhoY[:, igy, :], kernel, mode="valid")
 
-            bdry.set_explicit_boundary_data(results[n][0], elem, ud, th, mpv)
+            bdry_c.set_ghost_cells(results[n][0], elem, ud, th, npf)
             pn = np.expand_dims(pn, 1)
             pn = np.repeat(pn, node.icy, axis=1)
 
             results[n][2].p2_nodes[1:-1, :, 1:-1] = pn
-            bdry.set_ghostnodes_p2(results[n][2].p2_nodes, node, ud)
+            bdry_n.set_ghost_nodes(results[n][2].p2_nodes, node, ud)
 
             pn = np.expand_dims(results[n][2].p2_nodes[:, igy, :], 1)
             results[n][2].p2_nodes[...] = np.repeat(pn[...], node.icy, axis=1)

src/pybella/data_assimilation/prepare.py

@@ -48,19 +48,19 @@ def initialise(sst):
         logging.info("Seeds used in generating initial ensemble spread = ", seeds)
         for n in range(sst.N):
             Sol0 = deepcopy(sst.Sol)
-            mpv0 = deepcopy(sst.mpv)
+            npf0 = deepcopy(sst.npf)
             Sol0 = sst.sol_init(
-                Sol0, mpv0, es.elem, es.node, es.th, es.ud, seed=seeds[n]
+                Sol0, npf0, es.elem, es.node, es.th, es.ud, seed=seeds[n]
             )
-            # sol_ens[n] = [Sol0, deepcopy(es.flux), mpv0, [-np.inf, es.step]]
+            # sol_ens[n] = [Sol0, deepcopy(es.flux), npf0, [-np.inf, es.step]]
             sol_ens.update_member(
-                es.elem, es.node, Sol0, mpv0, deepcopy(es.flux), es.th
+                es.elem, es.node, Sol0, npf0, deepcopy(es.flux), es.th
             )
 
             sst.ensembble_state = sol_ens
     # elif sst.restart == False:
-    # sol_ens = [[sst.sol_init(mp.Sol, mp.mpv, mp.elem, mp.node, mp.th, sst.ud), mp.flux, mp.mpv, [-np.inf, sst.step]]]
-    # sol_ens.update_member(mp.elem, mp.node, sst.sol_init(mp.Sol, mp.mpv, mp.elem, mp.node, mp.th, sst.ud), mp.mpv, deepcopy(mp.flux), mp.th)
+    # sol_ens = [[sst.sol_init(mp.Sol, mp.npf, mp.elem, mp.node, mp.th, sst.ud), mp.flux, mp.npf, [-np.inf, sst.step]]]
+    # sol_ens.update_member(mp.elem, mp.node, sst.sol_init(mp.Sol, mp.npf, mp.elem, mp.node, mp.th, sst.ud), mp.npf, deepcopy(mp.flux), mp.th)
     # for n in range(sst.N):
     #     sol_ens.get_member(n).time.t = -np.inf
 

src/pybella/data_assimilation/utils.py

@@ -6,8 +6,7 @@
 import matplotlib.pyplot as plt
 
 from ..utils import options as opts
-
-from ..flow_solver.utils import boundary as bdry
+from ..flow_solver.utils.boundary import node_boundary as bdry_n
 
 
 class ensemble(object):
@@ -83,12 +82,12 @@ def set_p2_nodes(analysis, results, N, th, node, ud, loc_c=0, loc_n=2):
         rhoY_n[1:-1, 1:-1] = (
             sp.signal.fftconvolve(rhoY, kernel, mode="valid") / kernel.sum()
         )
-        bdry.set_ghostnodes_p2(rhoY_n, node, ud)
+        bdry_n.set_ghost_nodes(rhoY_n, node, ud)
         p2_n = rhoY_n**th.gm1 - 1.0 + (p2_n - p2_n.mean())
         # p2_n = rhoY_n**th.gm1 - 1.0
         p2_n -= p2_n.mean()
         # p2_n = np.pad(p2_n,2,mode='wrap')
-        bdry.set_ghostnodes_p2(p2_n, node, ud)
+        bdry_n.set_ghost_nodes(p2_n, node, ud)
         setattr(results[n][loc_n], "p2_nodes", p2_n)
 
 
@@ -260,7 +259,7 @@ def HSprojector_3t2D(results, elem, dap, N):
     Parameters
     ----------
     results : nd.array
-        An array of ensemble size k. Each ensemble member has [Sol,flux,mpv,[window_step,step]].
+        An array of ensemble size k. Each ensemble member has [Sol,flux,npf,[window_step,step]].
     dap : data assimilation input class
         .
     N : int