Merge pull request #80 from omltcat/xVel

Non-uniform x-velocity profile

docs/VnVPlan/VnVPlan.tex

@@ -739,6 +739,7 @@ \subsubsection{Flow Field Module [MG: M5]} \label{UT:Field}
   \item test-eddy-generation
   \item test-flow-field
   \item test-flow-field-wrap
+  \item test-flow-field-non-uniform-x-vel
   \item test-flow-field-init-exceptions
   \item test-flow-field-mesh-exceptions
   \item test-flow-field-set-exceptions

src/main.py

@@ -43,6 +43,12 @@ def main(args=None):
         type=float,
         help="Average flow velocity in x direction (default: 0)",
     )
+    new_parser.add_argument(
+        "-x",
+        default="",
+        metavar="X_FUNC",
+        help="X velocity profile function to be used (default: none)"
+    )
 
     # Query field subparser
     query_parser = subparsers.add_parser(
@@ -84,7 +90,7 @@ def main(args=None):
         try:
             profile = EddyProfile(args.p)
             field = FlowField(
-                profile=profile, name=args.n, dimensions=args.d, avg_vel=args.v
+                profile=profile, name=args.n, dimensions=args.d, avg_vel=args.v, x_vel_prof=args.x
             )
             field.save()
             print(f"New field '{args.n}' created and saved successfully")

src/modules/flow_field.py

@@ -11,6 +11,7 @@
 from modules import shape_function
 from modules import eddy
 from modules.eddy_profile import EddyProfile
+from modules import x_velocity
 
 WRAP_ITER = [-1, 0, 1]  # Iterations to wrap around the flow field, do not change
 CUTOFF = 1.2 * shape_function.get_cutoff()  # has to be greater than 1
@@ -32,6 +33,7 @@ def __init__(
         name: str,
         dimensions: np.ndarray | list,
         avg_vel: float | int = 0,
+        x_vel_prof: str = "",
     ):
         """
         Generate a new flow field.
@@ -59,6 +61,8 @@ def __init__(
             Dimensions of the the form of `[x, y, z]`
         `avg_vel` : float, optional (default: `0.0`)
             Average flow velocity to move the eddies
+        `x_vel_prof` : str, optional (default: `""`)
+            Name of the x-velocity profile to use, must already be defined in the `x_velocity.py`
         """
         if isinstance(dimensions, list):
             dimensions = np.array(dimensions)
@@ -115,7 +119,17 @@ def __init__(
             self.z[1] = self.z[0]
             self.y[2] = self.y[0]
             self.z[2] = self.z[0]
-        # if avg_vel is not zero, wrap around in x will have random y and z to avoid periodicity
+        # if avg_vel is not zero and velocity profile defined, use calculate individual x-velocity for each eddy
+        elif x_vel_prof != "":
+            self.x_vel_func = x_velocity.get_func(x_vel_prof)
+            print("Using x-velocity profile: ", x_vel_prof)
+            ny: np.ndarray = self.y[0] / self.high_bounds[1]
+            nz: np.ndarray = self.z[0] / self.high_bounds[2]
+            self.x_vel: np.ndarray = self.x_vel_func(ny, nz) * self.avg_vel
+            print("Max eddy center x-velocity: ", np.max(self.x_vel))
+            print("Min eddy center x-velocity: ", np.min(self.x_vel))
+        # if avg_vel is not zero and not velocity profile defined (constant avg_vel across entire field),
+        # wrap around in x will have random y and z to avoid periodicity
         else:
             self.set_rand_eddy_yz(1)
             self.set_rand_eddy_yz(2)
@@ -134,6 +148,11 @@ def get_eddy_centers(self, fi: int):
             self.set_rand_eddy_yz(fi)
         return np.stack((self.init_x, self.y[fi], self.z[fi]), axis=-1)
 
+    def get_eddy_center_x_vel(self, t):
+        """Get the x, y, and z coordinates of the eddies when x_vel of each eddy is defined."""
+        x = ((self.init_x + self.x_vel * t - self.low_bounds[0]) % self.dimensions[0]) + self.low_bounds[0]
+        return np.stack((x, self.y[0], self.z[0]), axis=-1)
+
     def set_rand_eddy_yz(self, fi: int):
         """Set random y and z coordinates for eddies in a new flow iteration."""
         self.y[fi] = np.random.uniform(self.low_bounds[1], self.high_bounds[1], self.N)
@@ -236,7 +255,25 @@ def sum_vel_mesh(
                     f"{e}\nNot enough memory to allocate velocity field. "
                     "Consider using a larger step size or focus on a smaller region."
                 ) from e
-            vel[..., 0] = self.avg_vel
+            if not hasattr(self, "x_vel_func"):
+                vel[..., 0] = self.avg_vel
+
+        # Initialize the x-velocity profile cross-section if needed
+        if hasattr(self, "x_vel_func"):
+            # Generate a plane of y and z coordinates
+            ny_coords = y_coords / self.high_bounds[1]
+            nz_coords = z_coords / self.high_bounds[2]
+            nynz = np.stack(
+                np.meshgrid(ny_coords, nz_coords, indexing="ij"), axis=-1
+            )[np.newaxis, ...]
+            ny = nynz[..., 0]
+            nz = nynz[..., 1]
+            # Calculate the x-velocity profile for each point on the plane
+            x_vel_plane = self.x_vel_func(ny, nz) * self.avg_vel
+            print("Max mean x-velocity: ", np.max(x_vel_plane))
+            print("Min mean x-velocity: ", np.min(x_vel_plane))
+        else:
+            x_vel_plane = None
 
         # Divide the coordinates into chunks
         if chunk_size == 0:     # pragma: no cover
@@ -269,7 +306,8 @@ def sum_vel_mesh(
         # Function to compute chunks looping through Y and Z for parallel processing of X
         def calc_x_chunks(i, xc):
             vel_i = np.zeros((len(xc), len(y_coords), len(z_coords), 3))
-            vel_i[..., 0] = self.avg_vel
+            if x_vel_plane is None:
+                vel_i[..., 0] = self.avg_vel
             mask = self.within_margin(
                 centers[:, 0], margins, x_coords[xc[0]], x_coords[xc[-1]]
             )
@@ -305,9 +343,10 @@ def calc_x_chunks(i, xc):
                         y_coords[yc],
                         z_coords[zc],
                     )
-            if self.verbose:
-                with plock:
-                    pbar.update(vel_i.size / 3)
+                    if self.verbose:
+                        pbar.update(1)
+            if x_vel_plane is not None:
+                vel_i[..., 0] += x_vel_plane
             if do_return:
                 vel[xc[0] : xc[-1] + 1, :, :, :] = vel_i
             if do_cache:
@@ -369,8 +408,12 @@ def get_wrap_arounds(
         # Wrap around for the x coordinates
         margin = self.sigma * CUTOFF
         for i in WRAP_ITER:
-            centers = self.get_eddy_centers(flow_iter + i)
-            centers[:, 0] += offset - i * self.dimensions[0]
+            if hasattr(self, "x_vel"):
+                centers = self.get_eddy_center_x_vel(t)
+                centers[:, 0] += i * self.dimensions[0]
+            else:
+                centers = self.get_eddy_centers(flow_iter + i)
+                centers[:, 0] += offset - i * self.dimensions[0]
             # Wrap around for the y and z coordinates
             for j in WRAP_ITER:
                 for k in WRAP_ITER:

src/modules/x_velocity.py

@@ -0,0 +1,43 @@
+"""
+Functions in this module is to calculate the multiplier of the mean x-velocity in the flow field.
+Arguments are the normalized distance from x-axis (centerline) in y and z directions.
+i.e. ny = 1 or -1 means the point is on the edge, ny = 0 means the point is on the centerline.
+ny and nz will always be in the range of -1 to 1.
+
+Also note that ny and nz are passed in as numpy arrays, so your function should be able to handle that.
+Generally, regular math operations will work just fine, but more complex logic may require additional considerations.
+
+If the function returns 0.5, and avg_vel set by user is 10 m/s, then the x-velocity at that point is 5 m/s.
+
+You can define your own velocity function here with any inner workings,
+as long as it takes ny and nz as arguments and returns a multiplier.
+"""
+
+from typing import Callable
+
+
+def get_func(func_name) -> Callable[[float, float], float]:
+    """
+    Get the function object by its name.
+    Don't change this function unless you know what you're doing.
+    """
+    try:
+        return globals()[func_name]
+    except KeyError:
+        raise ValueError(f"Velocity function \"{func_name}\" is not defined.")
+
+
+# Put your custom velocity functions below
+
+def parabola_2d(ny, nz):
+    """
+    Parabolic velocity profile in y, with no variation in z direction.
+    """
+    return 1 - ny**2
+
+
+def linear_2d(ny, nz):
+    """
+    Linear velocity profile in y, with no variation in z direction.
+    """
+    return (ny + 1) / 2

test.py

@@ -0,0 +1,11 @@
+import numpy as np
+cube = np.zeros((3, 3, 3, 3))
+plane = np.zeros((3, 3))
+plane += 1
+
+print(cube)
+
+# add plane to each slice of cube
+cube[..., 0] += plane
+
+print(cube[..., 0])

test/test_flow_field.py

@@ -289,6 +289,110 @@ def test_flow_field_parallel():
     assert len([f for f in os.listdir("src/.cache") if "x_" in f]) == expected_files
 
 
+@pytest.mark.unit
+def test_non_uniform_x_vel():
+    """Test using a linear velocity profile from 0 to 8 m/s from y=-10 to y=10 (like a slip boundary)"""
+    profile_name = "__test__"
+    content = {
+        "settings": {},
+        "variants": [
+            {"density": 0.005, "intensity": 1.2, "length_scale": 2.0},
+        ],
+    }
+
+    file_io.write("profiles", profile_name, content)
+    profile = EddyProfile(profile_name)
+
+    # Field creation
+    field_name = "two_eddies"
+    dimensions = np.array([20, 20, 20])
+    field = FlowField(profile, field_name, dimensions, avg_vel=8, x_vel_prof="linear_2d")
+
+    # Keep only two eddies, so we can track their movements without interference
+    field.N = 2
+    field.sigma = field.sigma[:2]
+    field.alpha = field.alpha[:2]
+
+    # Move one eddy to 0,-5,0 and the other to 0,5,0
+    field.init_x = np.array([0, 0])
+    field.y[0] = np.array([-5, 5])
+    field.z[0] = np.array([0, 0])
+
+    # Recalculate x_vel at these new positions
+    ny: np.ndarray = field.y[0] / field.high_bounds[1]
+    nz: np.ndarray = field.z[0] / field.high_bounds[2]
+    field.x_vel = field.x_vel_func(ny, nz) * field.avg_vel
+
+    field.save()
+
+    # Avg velocity for eddy A is 2 m/s, for eddy B is 6 m/s
+
+    # Get velocities at t=0, around 0,-5,0 and 0,5,0
+    vel_t0_a = field.sum_vel_mesh(
+        step_size=0.2,
+        chunk_size=5,
+        low_bounds=[-2, -7, -2],
+        high_bounds=[2, -3, 2],
+        time=0,
+    )
+    vel_t0_b = field.sum_vel_mesh(
+        step_size=0.2,
+        chunk_size=5,
+        low_bounds=[-2, 3, -2],
+        high_bounds=[2, 7, 2],
+        time=0,
+    )
+
+    # Get velocities at t=1, around 2,-5,0 and 6,5,0
+    vel_t1_a = field.sum_vel_mesh(
+        step_size=0.2,
+        chunk_size=5,
+        low_bounds=[0, -7, -2],
+        high_bounds=[4, -3, 2],
+        time=1,
+    )
+    vel_t1_b = field.sum_vel_mesh(
+        step_size=0.2,
+        chunk_size=5,
+        low_bounds=[4, 3, -2],
+        high_bounds=[8.0001, 7, 2],
+        time=1,
+    )
+    # Check if the velocity box at t=0 and t=1 is the same
+    diff_sum = np.sum(np.linalg.norm(vel_t0_a - vel_t1_a, axis=-1))
+    assert diff_sum < RTOL
+
+    diff_sum = np.sum(np.linalg.norm(vel_t0_b - vel_t1_b, axis=-1))
+    assert diff_sum < RTOL
+
+    # Get velocities at t=2, around 4,-5,0 and -8,5,0 (wrapping around)
+    vel_t2_a = field.sum_vel_mesh(
+        step_size=0.2,
+        chunk_size=5,
+        low_bounds=[2, -7, -2],
+        high_bounds=[6.0001, -3, 2],
+        time=2,
+    )
+    vel_t2_b = field.sum_vel_mesh(
+        step_size=0.2,
+        chunk_size=5,
+        low_bounds=[-10, 3, -2],
+        high_bounds=[-6, 7, 2],
+        time=2,
+    )
+
+    # Check if the velocity box at t=0 and t=2 is the same
+    diff_sum = np.sum(np.linalg.norm(vel_t0_a - vel_t2_a, axis=-1))
+    assert diff_sum < RTOL
+
+    diff_sum = np.sum(np.linalg.norm(vel_t0_b - vel_t2_b, axis=-1))
+    assert diff_sum < RTOL
+
+    # Clean up
+    os.remove(f"src/profiles/{profile_name}.json")
+    os.remove(f"src/fields/{field_name}.pkl")
+
+
 @pytest.mark.unit
 def test_flow_field_init_exceptions():
     """Test exceptions for flow field initialization"""

test/test_main.py

@@ -80,6 +80,46 @@ def test_main_query():
     assert glob.glob("src/results/main_test_*.npy")
 
 
+@pytest.mark.unit
+def test_main_new_x_profile():
+    """Test creating a new field with a custom x-velocity profile"""
+    args = [
+        "new",
+        "-p",
+        "main_test_profile",
+        "-n",
+        "main_test_field_x_profile",
+        "-d",
+        "10",
+        "10",
+        "10",
+        "-v",
+        "5",
+        "-x",
+        "parabola_2d",
+    ]
+    main.main(args)
+    assert os.path.exists("src/fields/main_test_field_x_profile.pkl")
+
+
+@pytest.mark.unit
+def test_main_query_x_profile():
+    """Test querying a field with a custom x-velocity profile"""
+    args = [
+        "query",
+        "-n",
+        "main_test_field_x_profile",
+        "-q",
+        "main_test_query",
+        "-s",
+        "gaussian",
+        "-c",
+        "1.0",
+    ]
+    main.main(args)
+    assert glob.glob("src/results/main_test_*.npy")
+
+
 @pytest.mark.unit
 def test_main_new_exceptions(capsys):
     """Test exceptions in creating a new field with main module as reported by lower level modules"""