initial commit

README.md

@@ -10,4 +10,11 @@ pinned: false
 license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+## Stacked Tensorial Neural Network (STNN) demo
+This demo uses the model architecture from [arXiv:2312.14979](https://arxiv.org/abs/2312.14979)
+to solve a parametric PDE problem on an elliptical annular domain. See the paper for a 
+detailed description of the problem and its applications.
+
+The [GitHub repo](https://github.com/caleb399/stacked_tensorial_nn) contains additional examples, including
+intructions for solving the PDE using a conventional iterative method (GMRES). Due to the long runtime of
+solving the PDE in this way, it is not included in the demo.

T5_config.json

@@ -0,0 +1,40 @@
+{
+    "K": 20,
+    "d": 8,
+    "W": 2,
+    "ranks": [
+        1,
+        16,
+        16,
+        16,
+        16,
+        16,
+        7,
+        1
+    ],
+    "shape1": [
+        4,
+        4,
+        4,
+        4,
+        4,
+        4,
+        4
+    ],
+    "shape2": [
+        4,
+        2,
+        2,
+        2,
+        2,
+        2,
+        2
+    ],
+    "nx1": 256,
+    "nx2": 64,
+    "nx3": 32,
+    "ell_min": 0.01,
+    "ell_max": 100.0,
+    "a2_min": 2.0,
+    "a2_max": 20.0
+}

T5_weights.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:51b60beefbbd6e5c8ade77d6f4f64c54146f34dc79e2b91605e0864ecbfcbd07
+size 957976

app.py

@@ -0,0 +1,316 @@
+import json
+import gradio as gr
+import matplotlib.pyplot as plt
+import matplotlib.ticker as ticker
+import numpy as np
+import sympy
+from matplotlib.cm import get_cmap
+from stnn.nn import stnn
+from stnn.pde.pde_system import PDESystem
+
+
+def adjust_to_nice_number(value, round_down = False):
+	"""
+	Adjust the given value to the nearest "nice" number. Used for colorbar tickmarks.
+	"""
+	if value == 0:
+		return value
+
+	is_negative = False
+	if value < 0:
+		round_down = True
+		is_negative = True
+		value = -value
+	exponent = np.floor(np.log10(value))  # Find exponent of 10
+	fractional_part = value / 10**exponent  # Find leading digit(s)
+
+	if round_down:
+		if fractional_part < 1.5:
+			nice_fractional = 1
+		elif fractional_part < 3:
+			nice_fractional = 2
+		elif fractional_part < 7:
+			nice_fractional = 5
+		else:
+			nice_fractional = 10
+	else:
+		if fractional_part <= 1:
+			nice_fractional = 1
+		elif fractional_part <= 2:
+			nice_fractional = 2
+		elif fractional_part <= 5:
+			nice_fractional = 5
+		else:
+			nice_fractional = 10
+
+	nice_value = nice_fractional * 10**exponent if round_down or nice_fractional != 10 else 10**(exponent + 1)
+	if is_negative:
+		nice_value = -nice_value
+	return nice_value
+
+
+def find_nice_values(min_val_raw, max_val, num_values = 4):
+	"""
+	Calculate 'num_values' evenly spaced "nice" values within the given range. Used for colorbar tickmarks.
+	"""
+	# Calculate rough spacing between values
+	min_val = adjust_to_nice_number(min_val_raw)
+	frac_val = (min_val - min_val_raw) / (max_val - min_val_raw)
+	if frac_val < 1 / num_values:
+		min_val = min_val_raw
+	raw_spacing = (max_val - min_val) / (num_values - 1)
+
+	# Calculate order of magnitude of the spacing
+	magnitude = np.floor(np.log10(raw_spacing))
+
+	nice_factors = np.array([1, 2, 5, 10])
+	normalized_spacing = raw_spacing / (10**magnitude)
+	closest_factor = nice_factors[np.argmin(np.abs(nice_factors - normalized_spacing))]
+	nice_spacing = closest_factor * (10**magnitude)
+
+	nice_values = min_val + nice_spacing * np.arange(num_values)
+
+	# Adjust if last value exceeds max_val
+	if nice_values[-1] < max_val - nice_spacing:
+		last_val = nice_values[-1]
+		nice_values = np.append(nice_values, [last_val + nice_spacing])
+
+	return [val for val in nice_values if min_val <= val <= max_val]
+
+
+def format_tick_label(val):
+	"""
+	Format w/ scientific notation for large/small values.
+	"""
+	if val != 0:
+		magnitude = np.abs(np.floor(np.log10(np.abs(val))))
+		if magnitude > 2:
+			return f'{val:.1e}'
+		elif magnitude > 1:
+			return f'{val:.0f}'
+		elif magnitude > 0:
+			return f'{val:.1f}'
+		else:
+			return f'{val:.2f}'
+	else:
+		return f'{val}'
+
+
+def plot_simple(system, rho, fontscale = 1):
+	# Major axis of outer boundary
+	b2 = system.b2
+
+	# Get x, y grids from 'PDESystem' object
+	x, y = system.get_xy_grids()
+
+	# wrap around values for continuity
+	rho = np.append(rho, rho[:, 0:1], axis = 1)
+
+	# Color bar limits
+	vmin = np.nanmin(rho)
+	vmax = np.nanmax(rho)
+
+	fig = plt.figure(figsize = (5, 5))
+	ax = plt.gca()
+
+	im = ax.contourf(x, y, rho, levels = np.linspace(vmin, vmax, 100), cmap = get_cmap('hsv'))
+	ax.set_title('rho(x,y)', fontsize = fontscale * 16)
+	for label in ax.get_xticklabels() + ax.get_yticklabels():
+		label.set_fontsize(fontscale * 12)
+	ax.set_aspect(1.0)
+	fac = 1.05
+	ax.set_xlim([-fac * b2, fac * b2])
+	ax.set_ylim([-fac * b2, fac * b2])
+
+	cbar = fig.colorbar(im, shrink = 0.8)
+
+	# Set colorbar ticks and labels to "nice" values
+	nice_values = find_nice_values(vmin, vmax, num_values = 5)
+	cbar.set_ticks(nice_values)
+	cbar.ax.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, pos: format_tick_label(x)))
+
+	return fig
+
+
+def evaluate_2d_expression(expr_str, xvals, yvals):
+	x, y = sympy.symbols('s t')
+	expr = sympy.sympify(expr_str)
+	f = sympy.lambdify((x, y), expr, modules = ['numpy'])
+	result = f(xvals, yvals)
+	if isinstance(result, (int, float)):
+		return result * np.ones(xvals.shape)
+	return f(xvals, yvals)
+
+
+'''
+# Currently unused in gradio interface
+def direct_solution(ell, a2, eccentricity, ibc_str, obc_str, max_krylov_dim, max_iterations):
+	# Direct solution
+	start = timeit.default_timer()
+
+	pde_config = {}
+	for key in ['nx1', 'nx2', 'nx3']:
+		pde_config[key] = stnn_config[key]
+	pde_config['ell'] = ell
+	pde_config['eccentricity'] = eccentricity
+	pde_config['a2'] = a2
+	system = PDESystem(pde_config)
+
+	try:
+		ibf_data = evaluate_2d_expression(ibc_str, system.x2_ib, system.x2_ib - system.x3_ib)[system.ib_slice]
+	except:
+		raise ValueError(f"Failed to parse the expression `{ibc_str}` for the boundary condition @ the inner boundary.")
+	try:
+		obf_data = evaluate_2d_expression(obc_str, system.x2_ob, system.x2_ob - system.x3_ob)[system.ob_slice]
+	except:
+		raise ValueError(f"Failed to parse the expression `{obc_str}` for the boundary condition @ the outer boundary.")
+
+	if np.any(np.isnan(ibf_data)):
+		raise ValueError(f"The expression `{ibc_str}` evaluates to nan at one or more grid points.")
+
+	if np.any(np.isnan(obf_data)):
+		raise ValueError(f"The expression `{obc_str}` evaluates to nan at one or more grid points.")
+
+	ibf_data, obf_data, b = system.convert_boundary_data(ibf_data, obf_data)
+
+	L_xp = csr_matrix(system.L)  # Sparse matrix representation of the PDE operator
+	nx1, nx2, nx3 = system.params['nx1'], system.params['nx2'], system.params['nx3']
+	b_xp = asarray(b.reshape((nx1 * nx2 * nx3,)))  # r.h.s. vector
+
+	def callback(res):
+		print(f'GMRES residual: {res}')
+
+	f_xp, info = spx.linalg.gmres(L_xp, b_xp, maxiter=max_iterations, tol=1e-7, restart=max_krylov_dim, callback=callback)
+
+	residual = (xp.linalg.norm(b_xp - L_xp @ f_xp) / xp.linalg.norm(b_xp))
+
+	if info > 0:
+		warnings.simplefilter('always')
+		warnings.warn(f'GMRES solver did not converge. Number of iterations: {info}; residual: {residual}', RuntimeWarning)
+
+	f = asnumpy(f_xp)
+	rho_direct = np.sum(f.reshape((nx1, nx2, nx3)), axis=-1)
+	direct_time = timeit.default_timer() - start
+	print(f'Done with direct solution. Time: {direct_time} seconds.')
+
+	fig = plot_simple(system, rho_direct)
+	return fig, info
+'''
+
+
+def predict_pde_solution(ell, a2, eccentricity, ibc_str, obc_str):
+	if a2 <= eccentricity:
+		raise ValueError(f'Outer minor axis must be greater than the eccentricity (here, {eccentricity}).')
+
+	pde_config = {}
+	for key in ['nx1', 'nx2', 'nx3']:
+		pde_config[key] = stnn_config[key]
+	pde_config['ell'] = ell
+	pde_config['eccentricity'] = eccentricity
+	pde_config['a2'] = a2
+	system = PDESystem(pde_config)
+
+	try:
+		ibf_data = evaluate_2d_expression(ibc_str, system.x2_ib, system.x2_ib - system.x3_ib)[system.ib_slice]
+	except:
+		raise ValueError(f"Failed to parse the expression `{ibc_str}` for the boundary condition @ the inner boundary.")
+	try:
+		obf_data = evaluate_2d_expression(obc_str, system.x2_ob, system.x2_ob - system.x3_ob)[system.ob_slice]
+	except:
+		raise ValueError(f"Failed to parse the expression `{obc_str}` for the boundary condition @ the outer boundary.")
+
+	if np.any(np.isnan(ibf_data)):
+		raise ValueError(f"The expression `{ibc_str}` evaluates to NaN at one or more grid points.")
+
+	if np.any(np.isnan(obf_data)):
+		raise ValueError(f"The expression `{obc_str}` evaluates to NaN at one or more grid points.")
+
+	# Permute and reshape boundary data to the format expected by the STNN model
+	ibf_data, obf_data, b = system.convert_boundary_data(ibf_data, obf_data)
+
+	'''
+	# Currently unused in gradio interface
+	ibf_data, obf_data, b, _ = system.generate_random_bc(func_gen_id)
+	'''
+
+	# Load some relevant quantities from the config dictionaries
+	ell_min, ell_max = stnn_config['ell_min'], stnn_config['ell_max']
+	a2_min, a2_max = stnn_config['a2_min'], stnn_config['a2_max']
+	nx1, nx2, nx3 = pde_config['nx1'], pde_config['nx2'], pde_config['nx3']
+
+	# Combine boundary data in single vector
+	bf = np.zeros((1, 2 * nx2, nx3 // 2))
+	bf[:, :nx2, :] = ibf_data[np.newaxis, ...]
+	bf[:, nx2:, :] = obf_data[np.newaxis, ...]
+
+	# Normalize and combine parameters
+	params = np.zeros((1, 3))
+	params[0, 0] = (a2 - a2_min) / (a2_max - a2_min)
+	params[0, 1] = (ell - ell_min) / (ell_max - ell_min)
+	params[0, 2] = eccentricity
+
+	rho = model.predict([params, bf])
+	fig = plot_simple(system, rho[0, ...])
+	return fig
+
+
+with open('T5_config.json', 'r', encoding = 'utf-8') as json_file:
+	stnn_config = json.load(json_file)
+
+model = stnn.build_stnn(stnn_config)
+model.load_weights('T5_weights.h5')
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Stacked Tensorial Neural Network (STNN) demo"
+                "\nThis demo uses the model architecture from [arXiv:2312.14979](https://arxiv.org/abs/2312.14979) "
+                "to solve a parametric PDE problem on an elliptical annular domain. "
+                "See the paper for a detailed description of the problem and its applications."
+                "<br/>The [GitHub repo](https://github.com/caleb399/stacked_tensorial_nn) contains additional examples, "
+                "including intructions for solving the PDE using a conventional iterative method (GMRES). "
+                "Due to the long runtime of solving the PDE in this way, it is not included in the demo.")
+    gr.Markdown("<br/>The PDE is "
+                "$\ell \\left( \\boldsymbol{\hat{u}} \cdot \\nabla \\right) f(\\boldsymbol{r}, w) = \partial_{ww} f(\\boldsymbol{r}, w)$, "
+                "where $\ell$ is a parameter and $\\boldsymbol{\hat{u}} = (\\cos w, \\sin w)$. "
+                "Here, $\\boldsymbol{r}$ is the 2D position vector, and $w$ is an angular coordinate unrelated to "
+                "the spatial domain. The model predicts the density !\\rho(\\boldsymbol{r}) = \int f(\\boldsymbol{r}, w) dw! "
+                "on elliptical annular domains parameterized as shown below. ",
+                latex_delimiters = [{"left": "$", "right": "$", "display": False}, {"left": "!", "right": "!", "display": True}])
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown(
+            "## PDE Parameters \n The model was trained on solutions of the PDE with $\ell$ between 0.01 and 100, $a$ between 2 and 20, "
+            "and $ecc$ between 0 and 0.8.", latex_delimiters = [{"left": "$", "right": "$", "display": False}, 
+            {"left": "!", "right": "!", "display": True}])
+            ell_input = gr.Number(label = "ell (must be > 0)", value = 1.0)
+            eccentricity_input = gr.Number(
+                label = "ecc: eccentricity of the inner boundary (must be >= 0 and <= 0.999)",
+                value = 0.5, minimum = 0.0, maximum = 0.999)
+            a2_input = gr.Number(label = "a: Minor axis of outer boundary (must be > eccentricity)", value = 2.0)
+            gr.Markdown(
+                "## Boundary Conditions \n $(s, t)$ are angular coordinates parameterizing the PDE domain, "
+                "related to $\\boldsymbol{r}$ and $w$ by a coordinate transformation. "
+                "Specifically, $s$ is the polar elliptical coordinate along the boundary (inner or outer), with values "
+                "between $-\pi$ and $\pi$, while $t = s - w$. Boundary conditions are generated from grid points "
+                "distributed uniformly over the allowable values of $s$ and $t$."
+                "<br/><br/>For the PDE problem to be well-posed, boundary data should only be specified where "
+                "$\\boldsymbol{\hat{u}} \cdot \\boldsymbol{\hat{n}} > 0$, where $\\boldsymbol{\hat{n}}$ is the "
+                "inward-pointing unit normal vector. This requirement constrains the allowable values of $t$."
+                " and is automatically enforced when building boundary conditions from the user-specified expressions below.",
+                latex_delimiters = [{"left": "$", "right": "$", "display": False}])
+
+            inner_boundary = gr.Textbox(label = "Inner boundary condition", value = "0.5 * (1 + sign(cos(s)))")
+            outer_boundary = gr.Textbox(label = "Outer boundary condition", value = "1 + 0.1 * cos(4*s)")
+
+            submit_button = gr.Button("Submit")
+
+        with gr.Column():
+            gr.Markdown("## Predicted Solution")
+            predicted_output_plot = gr.Plot()
+
+    submit_button.click(
+        fn = predict_pde_solution,
+        inputs = [ell_input, a2_input, eccentricity_input, inner_boundary, outer_boundary],
+        outputs = [predicted_output_plot]
+    )
+
+demo.launch()

requirements.txt

@@ -0,0 +1,10 @@
+tensorflow>=2.15.0
+numpy~=1.26.0
+t3f~=1.2.0
+scipy~=1.12.0
+h5py~=3.10.0
+matplotlib~=3.8.2
+pydot~=1.4.2
+openvino~=2023.3.0
+pyyaml>=6.0.1
+sympy

stnn/data/function_generators.py

@@ -0,0 +1,239 @@
+import numpy as np
+
+
+def generate_piecewise_linear_function(num_pieces, lower, upper, delta = 0.3):
+	"""
+	Generates a piece-wise linear function on the interval (lower, upper) that is 2*pi-periodic.
+
+	Args:
+		num_pieces (int): Number of linear pieces in the function.
+		lower (float): The lower range of the interval
+		upper (float): The upper range of the interval
+		delta (float): Parameter determining how rapidly the function varies between the grid points (i.e., modulates
+		the slopes of the piecewise functions). Larger values mean more variability. Default is 0.3.
+
+	Returns:
+		function: A piece-wise linear function.
+	"""
+	# Generate equally spaced points in the interval (-pi, pi)
+	x_points = np.linspace(lower, upper, num_pieces + 1)
+
+	# Generate random y-values for each point
+	y_points = np.zeros(num_pieces + 1)
+	y_points[0] = np.random.uniform(-1, 1)
+	for n in range(1, y_points.shape[0]):
+		y_points[n] = y_points[n - 1] + 0.3 * np.random.uniform(-1, 1)
+	y_points[0] = y_points[-1]
+	min_y = y_points.min()
+	# ensure y values are nonegative
+	if min_y < 0:
+		y_points -= min_y
+		y_points += np.random.uniform(0, 0.5)  # random (constant) offset
+
+	def piecewise_linear(x):
+		"""
+		Evaluates the piece-wise linear function at a given x.
+
+		Args:
+			x (float): The x-coordinate at which to evaluate the function.
+
+		Returns:
+			float: The y-coordinate of the function at x.
+		"""
+		for i in range(num_pieces):
+			if x_points[i] <= x < x_points[i + 1]:
+				# Linear interpolation between the two points
+				slope = (y_points[i + 1] - y_points[i]) / (x_points[i + 1] - x_points[i])
+				return slope * (x - x_points[i]) + y_points[i]
+		return y_points[0]  # For x = pi
+
+	return piecewise_linear
+
+
+def generate_piecewise_constant_function(num_pieces, lower, upper):
+	"""
+	Generates a piece-wise constant function on the interval (lower, upper) that is 2*pi periodic.
+
+	Args:
+		num_pieces (int): Number of constant pieces in the function.
+		lower (float): The lower range of the interval
+		upper (float): The upper range of the interval
+
+	Returns:
+		function: A piece-wise constant function.
+	"""
+	# Generate equally spaced points in the interval (-pi, pi)
+	x_points = np.linspace(lower, upper, num_pieces + 1)
+
+	# Generate random y-values for each constant piece
+	y_values = np.random.rand(num_pieces) * 2 - 0  # Random values between 0 and 1
+
+	# Ensure the function is 2*pi periodic
+	y_values = np.append(y_values, y_values[0])
+
+	def piecewise_constant(x):
+		"""
+		Evaluates the piece-wise constant function at a given x.
+
+		Args:
+			x (float): The x-coordinate at which to evaluate the function.
+
+		Returns:
+			float: The y-coordinate of the function at x.
+		"""
+		for i in range(num_pieces):
+			if x_points[i] <= x < x_points[i + 1]:
+				return y_values[i]
+		return y_values[0]  # For x = pi
+
+	return piecewise_constant
+
+
+def generate_piecewise_bc(x2_grid, x3_grid, num_pieces):
+	"""
+	Generates a piecewise linear function on the domain defined by x2_grid and x3_grid.
+
+	Args:
+		x2_grid (numpy.ndarray): A 2D array, x2 grid values
+		x3_grid (numpy.ndarray): A 2D array, x3 grid values
+		num_pieces (int): The number of pieces in the piecewise linear function.
+
+	Returns:
+		numpy.ndarray: A 2D array representing the piecewise linear function
+	"""
+	x2_fun = generate_piecewise_linear_function(num_pieces, lower = x2_grid.min(), upper = x2_grid.max())
+	x3_fun = generate_piecewise_linear_function(num_pieces, lower = x3_grid.min(), upper = x3_grid.max())
+
+	x2vals_1d = np.zeros_like(x2_grid[:, 0])
+	x3vals_1d = np.zeros_like(x3_grid[0, :])
+	for i in range(x2vals_1d.shape[0]):
+		x2vals_1d[i] = x2_fun(x2_grid[i, 0])
+	for i in range(x3vals_1d.shape[0]):
+		x3vals_1d[i] = x3_fun(x3_grid[0, i])
+	x2vals_2d, x3vals_2d = np.meshgrid(x2vals_1d, x3vals_1d, indexing = 'ij')
+	return x2vals_2d * x3vals_2d
+
+
+def random_2d_gaussian(theta, phi):
+	"""
+	Generates a 2D Gaussian G(x,y), where
+		x = np.cos(0.5 * freq_x * theta - phase_x)
+		y = np.cos(0.5 * freq_y * phi - phase_y)
+	Here, the frequencies and phases are randomly sampled, and (theta, phi) define a 2D meshgrid.
+
+	Args:
+		theta (numpy.ndarray): 2D array, meshgrid of the first coordinate
+		phi (numpy.ndarray): 2D array, meshgrid of the second coordinate
+
+	Returns:
+		numpy.ndarray: A 2D array representing the values of the Gaussian on the grid.
+	"""
+	phase_x = np.random.uniform(0, 2 * np.pi)
+	phase_y = np.random.uniform(0, 2 * np.pi)
+	freq_x = np.random.randint(1, 2)
+	freq_y = np.random.randint(1, 2)
+
+	x = np.cos(0.5 * freq_x * theta - phase_x)
+	y = np.cos(0.5 * freq_y * phi - phase_y)
+
+	sigma_x = np.random.uniform(0.1, 3.0)
+	sigma_y = np.random.uniform(0.1, 1.0)
+	rho = 0
+
+	covariance_matrix = np.array([[sigma_x**2, rho * sigma_x * sigma_y],
+								  [rho * sigma_x * sigma_y, sigma_y**2]])
+	inv_sigma_xx = 1.0 / sigma_x**2
+	inv_sigma_yy = 1.0 / sigma_y**2
+	inv_sigma_xy = -rho / (sigma_x * sigma_y)
+
+	if np.any(np.linalg.eigvals(covariance_matrix) < 0):
+		raise ValueError('Covariance matrix is not positive semi-definite.')
+
+	def gaussian_2d(x, y):
+		return np.exp(-0.5 * (inv_sigma_xx * x**2 + inv_sigma_yy * y**2 + 2 * inv_sigma_xy * x * y))
+
+	gaussian_values = gaussian_2d(x, y)
+	return gaussian_values
+
+
+def generate_random_functions(N, X, Y, num_terms = 16, min_freq = 1, max_freq = 16, func_gen_id = 0):
+	"""
+	Generates N random 2pi-periodic functions on a 2D grid as a Fourier series, with different types of
+	modulation applied to the amplitudes.
+
+	Args:
+		N (int): Number of functions to generate.
+		X (numpy.ndarray): 2D array representing the values of the first coordinate on the grid
+		Y (numpy.ndarray): 2D array representing the values of the second coordinate on the grid
+		num_terms (int, optional): Number of terms in the Fourier series expansion. Default is 16.
+		min_freq (int, optional): Minimum frequency for the Fourier series terms. Default is 1.
+		max_freq (int, optional): Maximum frequency for the Fourier series terms. Default is 16.
+		func_gen_id (int, optional): Type of function to generate based on the decay of the expansion coefficients
+									 as frequency is increased. Values can range from -1 to 4. Default is 0.
+
+	Returns:
+		numpy.ndarray: A 3D numpy array of shape (N, nx, ny) containing the function values.
+
+	Raises:
+		ValueError: If max_freq is less than min_freq or if an invalid func_gen_id is provided.
+	"""
+
+	# Check if the maximum frequency is less than the minimum frequency
+	if max_freq < min_freq:
+		raise ValueError('max_freq cannot be less than min_freq')
+
+	# Generate uniformly distributed functions if func_gen_id is -1
+	if func_gen_id == -1:
+		F_batch = np.random.uniform(0, 1, size = (N,) + X.shape)
+		return F_batch
+
+	# Initialize the batch of functions with zeros
+	F_batch = np.zeros((N,) + X.shape)
+
+	# Loop through each function to be generated
+	for n in range(N):
+		# Add a cosine term with a half frequency with 20% chance
+		if np.random.uniform(0, 1) < 0.2:
+			amp_cos_half = np.random.uniform(0, 1)  # Amplitude for cosine term
+			phase_cos_half = np.random.uniform(0, 2 * np.pi)  # Phase shift for cosine term
+			F_batch[n] += amp_cos_half * np.cos(0.5 * X - phase_cos_half)
+
+		# Fourier series
+		for _ in range(num_terms):
+			amplitude = np.random.uniform(-1, 1)  # Random amplitude for y-component
+			kx, ky = np.random.randint(min_freq, max_freq + 1, 2)  # Frequencies for x and y components
+			phase_x = np.random.uniform(0, 2 * np.pi)  # Phase shift for x-component
+			phase_y = np.random.uniform(0, 2 * np.pi)  # Phase shift for y-component
+
+			# Determine the coefficient amplitude based on the func_gen_id
+			if func_gen_id == 0:
+				# No decay applied to amplitude
+				pass
+			elif func_gen_id == 1:
+				if np.random.uniform(0, 1) < 0.5:
+					amplitude = amplitude / kx
+				else:
+					amplitude = amplitude / ky
+			elif func_gen_id == 2:
+				amplitude = amplitude / (kx * ky)
+			elif func_gen_id == 3:
+				amplitude = amplitude / (kx * kx * ky * ky)
+			elif func_gen_id == 4:
+				# Gaussian decay with random covariance matrix
+				sxx = np.random.uniform(0.1, 1.0)
+				syy = np.random.uniform(0.1, 1.0)
+				sxy = np.random.uniform(0.1, 1.0)
+				amplitude = amplitude * np.exp(-(sxx * kx**2 + syy * ky**2 + sxy * kx * ky))
+			else:
+				raise ValueError(
+					f'Invalid func_gen_id. Should be an integer in the range [-1, 4], but received {func_gen_id}')
+
+			# Add the term to the nth function in the batch
+			F_batch[n] += amplitude * np.cos(kx * X - phase_x) * np.cos(ky * Y - phase_y)
+
+		# Adjust the function to ensure it's positive
+		minF = np.min(F_batch[n])
+		if minF < 0:
+			F_batch[n] -= minF
+
+	return F_batch

stnn/data/preprocessing.py

@@ -0,0 +1,312 @@
+import h5py
+import numpy as np
+
+# If STRICT_WARNING = True, the program exits when negative values are detected in ibf, obf, or rho
+# This is important to check because negative values are unphysical.
+STRICT_WARNING = True
+
+
+def verify_nonnegative(fname, ibf, obf, rho):
+	"""
+	Check ibf, obf, and rho for negative values.
+	"""
+	found_warning = False
+	if np.any(ibf < 0):
+		print(f'Warning: negative values detected in array "ibf" in {fname}; min val: {ibf.min()}')
+		found_warning = True
+	elif np.any(obf < 0):
+		print(f'Warning: negative values detected in array "obf" in {fname}')
+		found_warning = True
+	elif np.any(rho < 0):
+		print(f'Warning: negative values detected in array "rho" in {fname}')
+		found_warning = True
+
+	if found_warning and STRICT_WARNING:
+		print(f'Exiting program. To avoid exiting on this warning, set STRICT_WARNING to False in {__file__.name}')
+		exit()
+
+
+def get_data_from_file(fname, nx2, nx3, Nrange = None):
+	"""
+	Retrieves training X from the given HDF5 file. Assumes that the PDE parameters
+	are stored in datasets with their respective names, i.e., 'ell', 'a1', 'a2'. Likewise,
+	the density rho(x1,x2) and boundary X ibf(x2,x3) / obf(x2,x3) are stored in datasets
+	'rho', 'ibf', and 'obf'.
+
+	Args:
+		 nx2 (int): Second grid dimension
+		 nx3 (int): Third grid dimension
+		 fname (str): Path to the HDF5 file containing the X.
+		 Nrange (tuple, optional): A tuple of two integers specifying the range of X to extract (start, end).
+		 						   Defaults to None.
+
+	Returns:
+		tuple: Tuple of extracted X
+
+	Raises:
+		ValueError: If the file does not contain the required datasets.
+	"""
+	if not isinstance(fname, str):
+		raise TypeError('Filename must be a string.')
+	type_check1 = not (Nrange is None or isinstance(Nrange, (tuple, list)))
+	type_check2 = False
+	if isinstance(Nrange, (tuple, list)):
+		type_check2 = len(Nrange) != 2
+		if not type_check2:
+			type_check2 = not all((isinstance(i, int) or i is None) for i in Nrange)
+	if type_check1 or type_check2:
+		raise TypeError('Nrange must be a length-2 tuple or list of integers.')
+
+	if Nrange is None:
+		N1, N2 = None, None
+	else:
+		N1, N2 = Nrange
+
+	# Check that all datasets are present
+	dset_names = ['ell', 'a1', 'a2', 'rho', 'ibf', 'obf']
+	with h5py.File(fname, 'r') as input_file:
+		missing_keys = [key for key in dset_names if key not in input_file.keys()]
+		if missing_keys:
+			raise ValueError(f"Missing / incorrectly labeled datasets in file {fname}.'"
+							 f"Could not find datasets: {', '.join(missing_keys)}")
+
+		ell = input_file['ell'][N1:N2]
+		a2 = input_file['a2'][N1:N2]  # minor axis of outer boundary
+		a1 = input_file['a1'][N1:N2]  # minor axis of inner boundary
+		eccentricity = np.ones_like(a1) - a1  # eccentricity of inner boundary
+		rho = input_file['rho'][N1:N2]
+		ibf = input_file['ibf'][N1:N2]  # boundary X on inner boundary
+		obf = input_file['obf'][N1:N2]  # boundary X on outer boundary
+
+		verify_nonnegative(fname, ibf, obf, rho)
+
+		# Combine 'ibf' and 'obf' into single array
+		N = rho.shape[0]
+		bf = np.zeros((N, 2 * nx2, nx3 // 2), dtype = np.float32)
+		bf[:, :nx2, :] = ibf
+		bf[:, nx2:, :] = obf
+
+	return a2, ell, eccentricity, bf, rho
+
+
+def reshape_and_stack(a2, ell, ecc):
+	a2 = a2.reshape((-1, 1))
+	ell = ell.reshape((-1, 1))
+	ecc = ecc.reshape((-1, 1))
+	return np.hstack([a2, ell, ecc])
+
+
+def apply_normalization(bf, rho):
+	fac = np.average(np.abs(rho), axis = (1, 2))
+	fac = fac.reshape((-1, 1, 1))
+	bf /= fac
+	rho /= fac
+	return bf, rho
+
+
+def load_data(files, nx2, nx3, ell_min, ell_max, a2_min, a2_max,
+			  Nrange_list = None, params_slice = None, normalize_data = False):
+	"""
+	Loads X from the specified files and processes it for use with the STNN.
+
+	Args:
+		 nx2 (int): Second grid dimension
+		 nx3 (int): Third grid dimension
+		 ell_min / ell_max (float): Minimum / maximum value of 'ell' over parameter space
+		 a2_min / a2_max (float): Minimum / maximum value of 'a2' over parameter space
+		 files (str or list of str): List of file paths containing the X
+		 Nrange_list (list of tuples, optional): Slice indices for the extracting X from the corresponding file. If
+		 										 given, must have the same number of elements as 'file_list'. Defaults
+		 										 to None.
+		 params_slice (slice, optional): Boolean array for selecting X over a subset of parameter space (ell, a1, a2).
+										 Defaults to None.
+		 normalize_data (bool, optional): Flag to normalize 'bf' and 'rho'. Defaults to False.
+
+	Returns:
+		tuple: A tuple containing the values of ell, a1, a2, bf, and rho. The parameters
+			   ell, a1, a2 are combined into a single array 'params'.
+	"""
+	if isinstance(files, (list, tuple)) and len(files) == 0:
+		raise ValueError(f'List of files provided to "load_data" is empty.')
+	if not isinstance(files, (list, tuple)):
+		files = [files]
+	if Nrange_list is None or len(Nrange_list) == 0:
+		# Default
+		Nrange_list = [None for _ in range(len(files))]
+	else:
+		# User-specified; check shapes
+		if not isinstance(Nrange_list, (list, tuple)):
+			Nrange_list = [Nrange_list]
+		if len(files) != len(Nrange_list):
+			raise ValueError('List of input files must have same length as list of Nrange tuples')
+	a2_list = []
+	ell_list = []
+	ecc_list = []
+	bf_list = []
+	rho_list = []
+
+	# Get X from each file and add to the lists
+	for file, Nrange in zip(files, Nrange_list):
+		a2, ell, ecc, bf, rho = get_data_from_file(file, nx2, nx3, Nrange = Nrange)
+		a2_list.append(a2)
+		ell_list.append(ell)
+		ecc_list.append(ecc)
+		bf_list.append(bf)
+		rho_list.append(rho)
+
+	a2 = np.concatenate(a2_list)
+	ell = np.concatenate(ell_list)
+	ecc = np.concatenate(ecc_list)
+	bf = np.vstack(bf_list)
+	rho = np.vstack(rho_list)
+
+	# Map ell and a2 values onto [0, 1]
+	ell = (ell - ell_min) / (ell_max - ell_min)
+	a2 = (a2 - a2_min) / (a2_max - a2_min)
+
+	params = reshape_and_stack(a2, ell, ecc)
+
+	if not params_slice is None:
+		# Extract subset of X, if params_slice is given
+		params = params[params_slice, ...]
+		bf = bf[params_slice, ...]
+		rho = rho[params_slice, ...]
+
+	if normalize_data:
+		bf, rho = apply_normalization(bf, rho)
+
+	return params, bf, rho
+
+
+def load_training_data(file_list, nx2, nx3, ell_min, ell_max, a2_min, a2_max, Nrange_list = None,
+					   params_slice = None, test_size = 0.1, random_state = 23, normalize_data = True):
+	"""
+	Loads training X from specified files and preprocesses it for use with training the STNN.
+
+	This function wraps the 'load_data' function, adding additional steps specific to preparing training X.
+
+	Args:
+		 nx2 (int): Second grid dimension
+		 nx3 (int): Third grid dimension
+		 ell_min / ell_max (float): Minimum / maximum value of 'ell' over parameter space
+		 a2_min / a2_max (float): Minimum / maximum value of 'a2' over parameter space
+		 file_list (list of str): List of file paths containing the X
+		 Nrange_list (list of tuples, optional): Slice indices for the extracting X from the corresponding file. If
+		 										 given, must have the same number of elements as 'file_list'. Defaults
+		 										 to None.
+		 params_slice (slice, optional): Boolean array for selecting X over a subset of parameter space (ell, a1, a2).
+										 Defaults to None.
+		 test_size (float, optional): Size of the test/validation dataset as a fraction of the total dataset size.
+		 							  Defaults to 0.1.
+		 random_state (int, optional): Random seed used to select the train-test split. Defaults to 23.
+		 normalize_data (bool, optional): Flag to normalize 'bf' and 'rho'. Defaults to False.
+
+	Returns:
+		tuple: A tuple containing the values of ell, a1, a2, bf, and rho. The parameters
+			   ell, a1, a2 are combined into a single array 'params'.
+"""
+	params, bf, rho = load_data(file_list, nx2, nx3, ell_min, ell_max, a2_min, a2_max,
+								Nrange_list = Nrange_list, params_slice = params_slice, normalize_data = normalize_data)
+
+	(rho_train, rho_test,
+	 Y_train, Y_test) = train_test_split(rho, [params, bf], test_size = test_size, random_state = random_state)
+
+	params_train = Y_train[0]
+	params_test = Y_test[0]
+	bf_train = Y_train[1]
+	bf_test = Y_test[1]
+
+	print('Finished loading training X:')
+	print(f'  params_train.shape:\t{params_train.shape}')
+	print(f'  bf_train.shape:\t{bf_train.shape}')
+	print(f'  rho_train.shape:\t{rho_train.shape}')
+	print(f'  params_test.shape:\t{params_test.shape}')
+	print(f'  bf_test.shape:\t{bf_test.shape}')
+	print(f'  rho_test.shape:\t{rho_test.shape}')
+
+	# Compute min/max extent of training X in parameter space.
+	# Note that 'params' is denormalized before computing the max/min.
+	min_a2 = np.min(a2_min + (a2_max - a2_min) * params[:, 0])
+	min_ell = np.min(ell_min + (ell_max - ell_min) * params[:, 1])
+	min_ecc = np.min(params[:, 2])
+
+	max_a2 = np.max(a2_min + (a2_max - a2_min) * params[:, 0])
+	max_ell = np.max(ell_min + (ell_max - ell_min) * params[:, 1])
+	max_ecc = np.max(params[:, 2])
+
+	print('')
+	print(f'  Number of circle samples  (train):\t{np.sum(params[:, 2] < 1e-7)}')
+	print(f'  Number of ellipse samples (train):\t{np.sum(params[:, 2] > 0)}')
+	print(f'  Min .. Max in training X:')
+	print(f'     ell:\t{min_ell:.2f} .. {max_ell:.2f}')
+	print(f'     a2:\t{min_a2:.2f} .. {max_a2:.2f}')
+	print(f'     ecc:\t{min_ecc:.2f} .. {max_ecc:.2f}')
+	print('-------------------------------------------')
+
+	return params_train, bf_train, rho_train, params_test, bf_test, rho_test
+
+
+def train_test_split(X, Y, test_size = 0.1, random_state = None):
+	"""
+	Split (X, Y) pairs into random train and test subsets.
+
+	Args:
+		X (np.ndarray or list of arrays): Input dataset
+		Y (np.ndarray or list of arrays): Labels for the dataset
+		test_size (float): Proportion of the dataset to include in the test split
+		random_state (int): Controls the shuffling applied to the X and Y before applying the split
+
+	Returns:
+		X_train, X_test, Y_train, Y_test: Lists containing train-test split of the dataset. The format is
+		the same as the input X. For example, if 'X' is an array and 'Y' is a list of arrays, then X_train
+		and X_test will be arrays, and Y_train and Y_test will be lists of arrays.
+
+	Note: This function is included primarilyto reduce module dependency requirements, and it may not be memory-efficient
+		  for large datasets. sklearn.model_selection.train_test_split has similar functionality and may be preferred
+		  for performance-critical applications.
+	"""
+	if len(X) == 0 or len(Y) == 0:
+		raise ValueError("Input arrays/lists X and Y cannot be empty.")
+
+	input_X_is_array = isinstance(X, np.ndarray)
+	input_Y_is_array = isinstance(Y, np.ndarray)
+
+	if input_X_is_array:
+		X = [X]
+	if input_Y_is_array:
+		Y = [Y]
+
+	total_samples = X[0].shape[0]
+
+	# Check for consistent number of samples across all datasets
+	if any(x.shape[0] != total_samples for x in X) or any(y.shape[0] != total_samples for y in Y):
+		raise ValueError('Inconsistent number of samples.')
+
+	Ntest = int(test_size * total_samples)
+	if Ntest < 1 or Ntest > total_samples:
+		raise ValueError('Size of test dataset cannot be less than 1 or greater than the total number of samples.')
+
+	if random_state is not None:
+		np.random.seed(random_state)
+
+	# Shuffle indices
+	indices = np.arange(total_samples)
+	np.random.shuffle(indices)
+
+	# Apply shuffled indices to all datasets
+	shuffled_X = [x[indices] for x in X]
+	shuffled_Y = [y[indices] for y in Y]
+
+	# Split X and Y
+	X_train = [x[:-Ntest] for x in shuffled_X]
+	X_test = [x[-Ntest:] for x in shuffled_X]
+	Y_train = [y[:-Ntest] for y in shuffled_Y]
+	Y_test = [y[-Ntest:] for y in shuffled_Y]
+
+	# Convert back to arrays if original input was array
+	if input_X_is_array:
+		X_train, X_test = X_train[0], X_test[0]
+	if input_Y_is_array:
+		Y_train, Y_test = Y_train[0], Y_test[0]
+
+	return X_train, X_test, Y_train, Y_test

stnn/data/test_functions.py

@@ -0,0 +1,197 @@
+import numpy as np
+
+"""
+Interface to a subset of the test functions listed at
+https://en.wikipedia.org/wiki/Test_functions_for_optimization
+"""
+
+
+def rastrigin(x, y):
+	args = (x, y)
+	A = 10
+	return A * len(x) + sum([(xi**2 - A * np.cos(2 * np.pi * xi)) for xi in args])
+
+
+def ackley(x, y):
+	return -20 * np.exp(-0.2 * np.sqrt(0.5 * (x**2 + y**2))) - \
+		np.exp(0.5 * (np.cos(2 * np.pi * x) + np.cos(2 * np.pi * y))) + np.e + 20
+
+
+def sphere(x, y):
+	return x**2 + y**2
+
+
+def rosenbrock(x, y):
+	return 100 * (y - x**2)**2 + (1 - x)**2
+
+
+def beale(x, y):
+	return (1.5 - x + x * y)**2 + (2.25 - x + x * y**2)**2 + (2.625 - x + x * y**3)**2
+
+
+def goldstein_price(x, y):
+	return (1 + (x + y + 1)**2 * (19 - 14 * x + 3 * x**2 - 14 * y + 6 * x * y + 3 * y**2)) * \
+		(30 + (2 * x - 3 * y)**2 * (18 - 32 * x + 12 * x**2 + 48 * y - 36 * x * y + 27 * y**2))
+
+
+def booth(x, y):
+	return (x + 2 * y - 7)**2 + (2 * x + y - 5)**2
+
+
+def bukin(x, y):
+	return 100 * np.sqrt(abs(y - 0.01 * x**2)) + 0.01 * abs(x + 10)
+
+
+def matyas(x, y):
+	return 0.26 * (x**2 + y**2) - 0.48 * x * y
+
+
+def levi(x, y):
+	return np.sin(3 * np.pi * x)**2 + (x - 1)**2 * (1 + np.sin(3 * np.pi * y)**2) + \
+		(y - 1)**2 * (1 + np.sin(2 * np.pi * y)**2)
+
+
+def himmelblau(x, y):
+	return (x**2 + y - 11)**2 + (x + y**2 - 7)**2
+
+
+def three_hump_camel(x, y):
+	return 2 * x**2 - 1.05 * x**4 + x**6 / 6 + x * y + y**2
+
+
+def easom(x, y):
+	return -np.cos(x) * np.cos(y) * np.exp(-((x - np.pi)**2 + (y - np.pi)**2))
+
+
+def cross_in_tray(x, y):
+	return -0.0001 * (abs(np.sin(x) * np.sin(y) * np.exp(abs(100 - np.sqrt(x**2 + y**2) / np.pi))) + 1)**0.1
+
+
+def eggholder(x, y):
+	return -(y + 47) * np.sin(np.sqrt(abs(x / 2 + (y + 47)))) - x * np.sin(np.sqrt(abs(x - (y + 47))))
+
+
+def holder_table(x, y):
+	return -abs(np.sin(x) * np.cos(y) * np.exp(abs(1 - np.sqrt(x**2 + y**2) / np.pi)))
+
+
+def mccormick(x, y):
+	return np.sin(x + y) + (x - y)**2 - 1.5 * x + 2.5 * y + 1
+
+
+def schaffer2(x, y):
+	return 0.5 + (np.sin(x**2 - y**2)**2 - 0.5) / (1 + 0.001 * (x**2 + y**2))**2
+
+
+def schaffer4(x, y):
+	return 0.5 + (np.cos(np.sin(abs(x**2 - y**2)))**2 - 0.5) / (1 + 0.001 * (x**2 + y**2))**2
+
+
+def styblinski_tang(x, y):
+	args = (x, y)
+	return sum([xi**4 - 16 * xi**2 + 5 * xi for xi in args]) / 2
+
+
+functions = [
+	rastrigin,
+	ackley,
+	sphere,
+	rosenbrock,
+	beale,
+	goldstein_price,
+	booth,
+	bukin,
+	matyas,
+	levi,
+	himmelblau,
+	three_hump_camel,
+	easom,
+	cross_in_tray,
+	eggholder,
+	holder_table,
+	mccormick,
+	schaffer2,
+	schaffer4,
+	styblinski_tang
+]
+
+function_names = [
+	'rastrigin',
+	'ackley',
+	'sphere',
+	'rosenbrock',
+	'beale',
+	'goldstein_price',
+	'booth',
+	'bukin',
+	'matyas',
+	'levi',
+	'himmelblau',
+	'three_hump_camel',
+	'easom',
+	'cross_in_tray',
+	'eggholder',
+	'holder_table',
+	'mccormick',
+	'schaffer2',
+	'schaffer4',
+	'styblinski_tang'
+]
+
+domains = {
+	'rastrigin': (-5.12, 5.12),
+	'ackley': (-5, 5),
+	'sphere': (-1, 1),
+	'rosenbrock': {'x': (-2, 2), 'y': (-10, 10)},
+	'beale': (-4.5, 4.5),
+	'goldstein_price': (-2, 2),
+	'booth': (-10, 10),
+	'bukin': {'x': (-15, -5), 'y': (-3, 3)},
+	'matyas': (-10, 10),
+	'levi': (-10, 10),
+	'himmelblau': (-5, 5),
+	'three_hump_camel': (-5, 5),
+	'easom': (-100, 100),
+	'cross_in_tray': (-10, 10),
+	'eggholder': (-512, 512),
+	'holder_table': (-10, 10),
+	'mccormick': {'x': (-1.5, 4), 'y': (-3, 4)},
+	'schaffer2': (-100, 100),
+	'schaffer4': (-100, 100),
+	'styblinski_tang': (-5, 5)
+}
+
+
+def scale_input(x, domain):
+	min_d, max_d = domain
+	return min_d + (max_d - min_d) * x
+
+
+def get_test_function(X, Y, fun_idx):
+	"""
+	Evaluates a function on inputs (X, Y).
+
+	Args:
+		X (float or array-like): The X input values to be scaled and used in the function.
+		Y (float or array-like): The Y input values to be scaled and used in the function.
+								Ignored if the function takes only one argument.
+		fun_idx (int): The index of the function to be retrieved from a predefined list 'functions'.
+
+	Returns: 
+		(float or array-like), values of the function on the grid
+	"""
+	func = functions[fun_idx]
+
+	domain = domains[func.__name__]
+	if isinstance(domain, dict):
+		x_scaled = scale_input(X, domain['x'])
+		y_scaled = scale_input(Y, domain['y'])
+	else:
+		x_scaled = scale_input(X, domain)
+		y_scaled = scale_input(Y, domain)
+
+	try:
+		output = func(X, Y)
+	except TypeError:
+		output = func(X)
+	return output

stnn/linalg_backend.py

@@ -0,0 +1,81 @@
+"""
+Module imports and wrapper functions of the linear algebra backend.
+
+If  __usecupy__ is "True" and cupy is successfully imported, then
+
+	xp 		--> cupy
+	spx		--> cupyx.scipy.sparse.linalg
+	
+Otherwise,
+
+	xp 		--> numpy
+	spx		--> scipy.sparse
+
+For example, if __usecupy__ is False, then
+
+				import numpy as np
+				import scipy.sparse.linalg as sp
+				
+is equivalent to
+
+				from stnn.linalg_backend import xp, spx
+"""
+__usecupy__ = True
+
+try:
+	# If CuPy is not preferred or available, fall back to NumPy
+	if not __usecupy__:
+		raise ImportError
+	import cupy as cp
+	import cupyx.scipy.sparse.linalg
+	import cupyx.scipy.sparse as cupy_sparse
+
+	xp = cp
+	spx = cupy_sparse
+	using_cupy = True
+except ImportError:
+	import numpy as np
+	import scipy.sparse.linalg
+	import scipy.sparse as scipy_sparse
+
+	xp = np
+	spx = scipy_sparse
+	using_cupy = False
+
+
+def csr_matrix(L):
+	"""
+	Create a CSR (Compressed Sparse Row) matrix.
+
+	If CuPy is available and enabled, this function will create a CuPy CSR matrix.
+	Otherwise, it converts the given data to a SciPy CSR matrix.
+
+	Parameters:
+	L (array_like or sparse matrix): 2-D array or sparse matrix to convert.
+
+	Returns:
+	CSR matrix: The converted CSR matrix, using either CuPy or SciPy.
+	"""
+	if using_cupy:
+		return spx.csr_matrix(L, dtype=xp.float64)
+	return L.tocsr()
+
+
+def asnumpy(arr):
+	"""
+	Convert an array from the backend library (CuPy or NumPy) to NumPy.
+	If NumPy is enabled, the input array is returned unchanged.
+	"""
+	if using_cupy:
+		return cp.asnumpy(arr)
+	return arr
+
+
+def asarray(arr):
+	"""
+	Convert the input to an array of the backend library (CuPy or NumPy).
+	If NumPy is enabled, the input array is returned unchanged.
+	"""
+	if using_cupy:
+		return cp.asarray(arr, dtype=cp.float64)
+	return arr

stnn/nn/stnn.py

@@ -0,0 +1,66 @@
+import tensorflow as tf
+from tensorflow.keras.layers import Input, Multiply, Add
+from tensorflow.keras.models import Model
+
+from .stnn_layers import TTL, SoftmaxEmbeddingLayer
+
+
+def build_stnn(config):
+	"""
+	Constructs a Stacked Tensorial Neural Network (STNN) as a TensorFlow model based on
+	the provided configuration dictionary.
+
+	Args:
+		config (dict): Configuration dictionary for the STNN model. Must contain the following entries:
+						- 'K' (int): Number of tensor networks to be stacked
+						- 'd' (int): Number of dense layers in the model's SoftmaxEmbeddingLayer.
+						- 'nx1', 'nx2', 'nx3' (int): Dimensions of the finite-difference grid
+						- All other required entries for the TTL class, not already listed above.
+
+	Returns:
+		tf.keras.Model: The constructed STNN model.
+
+	Raises:
+		ValueError: If the config dictionary does not contain positive integers 'K', 'd', 'nx2', 'nx3';
+					also if config['nx3'] is not divisible by 2.
+	"""
+	required_keys = ['nx1', 'nx2', 'nx3', 'K', 'd', 'shape1','shape2','ranks','W']
+	missing_keys = [key for key in required_keys if key not in config]
+	if missing_keys:
+		raise KeyError(f"Missing keys in config: {', '.join(missing_keys)}")
+
+	for key in ['nx1', 'nx2', 'nx3', 'K', 'd']:
+		if not isinstance(config[key], int):
+			raise TypeError(f"{key} must be an integer.")
+
+	for key in ['nx1', 'nx2', 'nx3', 'K', 'd']:
+		if config[key] <= 0:
+			raise ValueError(f"{key} must be positive.")
+
+	if config['nx3'] % 2 == 1:
+		raise ValueError('Config error: nx3 must be divisible by 2.')
+
+	K = config['K'] # Number of tensor networks
+	d = config['d'] # Number of dense layers in SoftmaxEmbeddingLayer
+	input_shape = (2 * config['nx2'], config['nx3'] // 2, 1)
+	input_tensor = Input(shape = input_shape)
+
+	# Process parameter array (ell, a1, a2) and output weights for stacking the tensor networks
+	preprocess_layer = SoftmaxEmbeddingLayer(K, d)
+	params_input = Input(shape = (3,))
+	stack_weights = preprocess_layer(params_input)[:, tf.newaxis, tf.newaxis, :]
+
+	# Build the tensor networks using the custom keras layer class TLL
+	models = [TTL(config) for _ in range(K)]
+
+	# Combine the tensor networks based on the weights outputted by 'preprocess_layer'
+	weighted_outputs = []
+	for i, model in enumerate(models):
+		processed_output = model(input_tensor)
+		weighted_output = Multiply()([processed_output, stack_weights[..., i]])
+		weighted_outputs.append(weighted_output)
+	final_output = Add()(weighted_outputs)
+
+	model = Model(inputs = [params_input, input_tensor], outputs = final_output)
+
+	return model

stnn/nn/stnn_layers.py

@@ -0,0 +1,274 @@
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.layers import Reshape, Flatten
+import t3f
+
+import os
+import logging
+
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'  # FATAL
+logging.getLogger('tensorflow').setLevel(logging.FATAL)
+
+
+class SoftmaxEmbeddingLayer(tf.keras.layers.Layer):
+	"""
+	Parameter embedding layer that generates the weights used for stacking the tensor networks. It
+	takes the parameter array, lambda = (ell, a1, a2), as input and outputs K numbers that sum to 1.
+
+	Attributes:
+		output_dim (int): The dimension of the output
+		expansion_dim (int): The dimension used for expanding the input in intermediate layers.
+	"""
+
+	def __init__(self, output_dim, d, expansion_dim = 30, **kwargs):
+		super(SoftmaxEmbeddingLayer, self).__init__(**kwargs)
+		self.reduction_layer = None
+		self.expansion_layers = None
+		self.output_dim = output_dim
+		self.expansion_dim = expansion_dim
+		self.d = d  # Number of dense layers
+
+	def build(self, input_shape):
+		# Expansion layers to increase dimensionality
+		self.expansion_layers = [tf.keras.layers.Dense(self.expansion_dim, activation = 'relu') for _ in range(self.d)]
+		# Reduction layer to bring dimensionality back to the desired output dimension
+		self.reduction_layer = tf.keras.layers.Dense(self.output_dim)
+
+	def call(self, inputs):
+		expanded = inputs
+		for layer in self.expansion_layers:
+			expanded = layer(expanded)
+		return tf.nn.softmax(self.reduction_layer(expanded))
+
+	def get_config(self):
+		return {'output_dim': self.output_dim, 'expansion_dim': self.expansion_dim}
+
+
+class EinsumTTLRegularizer(tf.keras.regularizers.Regularizer):
+	"""
+	Regularizer for the Einsum layer of the TTL layer class, penalizing high-frequency components of the
+	weights vector.
+
+	Attributes:
+		strength (float): The regularization strength.
+		midpoint (int): Index demarcating the inner and outer boundaries, i.e. x[:midpoint] contains
+						data for the inner boundary, and x[midpoint:] contains data for the outer boundary.
+						The regularization is designed so it does not penalize variations across this index.
+	"""
+
+	def __init__(self, strength, midpoint):
+		self.strength = strength
+		self.midpoint = midpoint
+
+	def __call__(self, x):
+		diff = tf.abs(x[1:self.midpoint - 1] - x[0:self.midpoint - 2]) \
+			   + tf.abs(x[self.midpoint + 1:2 * self.midpoint - 1] - x[self.midpoint:2 * self.midpoint - 2])
+		return self.strength * tf.reduce_sum(diff)
+
+	def get_config(self):
+		return {'strength': self.strength, 'midpoint': self.midpoint}
+
+
+def cosine_initializer(kx = 1.0):
+	"""
+	Initializer for the Einsum layer of the TTL layer class. Sets the weights to a linear combination
+	of cos(kx * x) and cos(2 * kx * x), where x is the weight vector.
+
+	Args:
+		kx (float, optional): Frequency of the cosine terms. Defaults to 1.0.
+
+	Returns:
+		_initializer: Weight initializer function
+	"""
+
+	def _initializer(shape, dtype = None):
+		x_values = np.linspace(-np.pi, np.pi, shape[0])
+		cos_values = np.random.uniform(-0.1, 0.3) * np.abs(np.cos(kx * x_values)) \
+					 + np.random.uniform(-0.05, 0.05) * np.abs(np.cos(2.0 * kx * x_values))
+		return tf.convert_to_tensor(-cos_values, dtype = dtype)
+
+	return _initializer
+
+
+class EinsumTTL(tf.keras.layers.Layer):
+	"""
+	Layer that contracts the input tensor over the second dimension before passing it to the TTL.
+	If regularization is enabled, it applies an `EinsumTTLRegularizer` to the kernels.
+
+	Attributes:
+	    (nx2, nx3) (integers): Shape parameters characterizing input tensor dimensions. T
+	                           The shape of the input tensor is (2*nx2, nx3//2).
+	    W (int): Number of einsum contractions
+	    kernels (list): List of weight matrices for each einsum contraction
+	    regularization_strength (float): The strength of the regularization if used.
+	    use_regularization (bool): Flag to indicate whether regularization is used.
+	"""
+
+	def __init__(self, nx2, nx3, W, use_regularization, regularization_strength = 0.005, **kwargs):
+		super(EinsumTTL, self).__init__(**kwargs)
+		self.nx2 = nx2
+		self.nx3 = nx3
+		self.W = W
+		self.kernels = []
+
+		self.regularization_strength = regularization_strength
+		self.use_regularization = use_regularization
+		if self.use_regularization:
+			regularizer = EinsumTTLRegularizer(self.regularization_strength, self.nx3 // 4)
+		else:
+			regularizer = None
+
+		initializer_values_ = [1.0, 0.5, 2.0, 3.0] * W
+		initializer_values = initializer_values_[:W]
+		for i in range(W):
+			self.kernels.append(self.add_weight(
+				name = f'w{i + 1}',
+				shape = (nx3 // 2,),
+				regularizer = regularizer,
+				initializer = cosine_initializer(initializer_values[i])
+			))
+
+	def call(self, inputs):
+		parts = []
+		for w in self.kernels:
+			part_a = tf.einsum('abc,c->ab', inputs[:, :self.nx2, :self.nx3 // 4], w[:self.nx3 // 4]) + \
+					 tf.einsum('abc,c->ab', inputs[:, :self.nx2, self.nx3 // 4:self.nx3 // 2],
+							   tf.reverse(w[:self.nx3 // 4], axis = [0]))
+			part_b = tf.einsum('abc,c->ab', inputs[:, self.nx2:, :self.nx3 // 4], w[self.nx3 // 4:self.nx3 // 2]) + \
+					 tf.einsum('abc,c->ab', inputs[:, self.nx2:, self.nx3 // 4:self.nx3 // 2],
+							   tf.reverse(w[self.nx3 // 4:self.nx3 // 2], axis = [0]))
+			parts.extend([part_a, part_b])
+
+		return tf.concat(parts, axis = 1)
+
+	def get_config(self):
+		return {'use_regularization': self.use_regularization,
+				'regularization_strength': self.regularization_strength}
+
+
+class TTL(tf.keras.layers.Layer):
+	"""
+	TTL (Tensor Train Layer) is a custom TensorFlow Keras layer that builds a model
+	based on the given configuration. This layer is designed to work with 
+	tensor train decomposition in neural networks.
+
+	Attributes:
+		config (dict): Configuration dictionary containing parameters for the model.
+		
+			'nx1', 'nx2', 'nx3': Integers, dimensions of the finite-difference grid
+			
+			'shape1': List of integers, defines the shape of the output tensor in the tensor train format.
+					  The length of shape1 must match the length of shape2.
+			'shape2': List of integers, specifies the shape of the input tensor in the tensor train format.
+					  The length of shape2 must match the length of shape1.
+			'ranks':  List of integers, represents the ranks in the tensor train decomposition.
+					  The length of this list determines the complexity and the number of parameters in the tensor train layer.
+			'W' (int): 	Number of weight vectors to use in the initial EinsumTTL layer. Setting W = 0 means that no EinsumTLL
+						used.
+
+			'use_regularization' (boolean, optional, default: False):  Indicates whether regularization is used in the EinsumTTL.
+			'regularization_strength' (float, optional, default: 0): Strength of the regularization
+			
+		model (tf.keras.Sequential): The Sequential model built based on the provided configuration.
+
+	Methods:
+		load_config(self, config): Loads configuration
+		build_model(self): Builds the layer
+		call(inputs): Method for the forward pass of the layer.
+	"""
+
+	def __init__(self, config, **kwargs):
+		super(TTL, self).__init__(**kwargs)
+		self.model = Sequential()
+		self.nx1 = None
+		self.nx2 = None
+		self.nx3 = None
+		self.shape1 = None
+		self.shape2 = None
+		self.ranks = None
+		self.W = None
+		self.use_regularization = None
+		self.regularization_strength = None
+		self._required_keys = ['nx1', 'nx2', 'nx3', 'shape1', 'shape2', 'ranks', 'W']
+
+		config.setdefault('use_regularization', False)
+		config.setdefault('regularization_strength', 0.0)
+		self.load_config(config)
+		self.config = config
+		self.build_model()
+
+	def load_config(self, config):
+		missing_keys = [key for key in self._required_keys if key not in config]
+		if missing_keys:
+			raise KeyError(f"Missing keys in config: {', '.join(missing_keys)}")
+
+		if not isinstance(config['use_regularization'], bool):
+			raise TypeError('use_regularization must be a boolean.')
+		else:
+			self.use_regularization = config['use_regularization']
+
+		self.regularization_strength = 0.0
+
+		for key in ['nx1', 'nx2', 'nx3', 'W']:
+			if not isinstance(config[key], int):
+				raise TypeError(f"{key} must be an integer.")
+
+		for key in ['nx1', 'nx2', 'nx3']:
+			if config[key] <= 0:
+				raise ValueError(f"{key} must be positive.")
+
+		if config['W'] < 0:
+			raise ValueError("W must be non-negative.")
+
+		nx1, nx2, nx3 = config['nx1'], config['nx2'], config['nx3']
+		self.nx1 = nx1
+		self.nx2 = nx2
+		self.nx3 = nx3
+
+		W = config['W']
+		self.W = W
+
+		input_dim = 2 * nx2 * W
+		if W == 0:
+			input_dim = nx2 * nx3
+
+		shape1, shape2 = config['shape1'], config['shape2']
+		if len(shape1) != len(shape2):
+			raise ValueError(
+				f'shape1 and shape2 must have the same length. '
+				f'Received: shape1 = {shape1}, shape2 = {shape2}.'
+			)
+		elif np.prod(np.array(shape1)) != nx1 * nx2:
+			raise ValueError(
+				f'prod(shape1) must be equal to the output dimension of the TTL '
+				f'(nx1 * nx2,). Received: prod(shape1) = {np.prod(np.array(shape1))}, '
+				f'nx1 * nx2 = {nx1 * nx2}.'
+			)
+		elif np.prod(np.array(shape2)) != input_dim:
+			raise ValueError(
+				f'prod(shape2) must be equal to the input dimension of the TTL '
+				f'(2 * nx2 * W or nx2 * nx3 if W = 0). '
+				f'Received: prod(shape2) = {np.prod(np.array(shape2))}, required input dimension = {input_dim}.'
+			)
+		else:
+			self.shape1 = shape1
+			self.shape2 = shape2
+
+		self.ranks = config['ranks']
+
+	def build_model(self):
+		if self.W == 0:
+			self.model.add(Flatten(input_shape = (2 * self.nx2, self.nx3 // 2)))
+		else:
+			self.model.add(EinsumTTL(self.nx2, self.nx3, self.W, self.use_regularization,
+									 regularization_strength = self.regularization_strength,
+									 input_shape = (2 * self.nx2, self.nx3 // 2)))
+			self.model.add(Flatten())
+		tt_layer = t3f.nn.KerasDense(input_dims = self.shape2, output_dims = self.shape1,
+									 tt_rank = self.ranks, use_bias = False, activation = 'linear')
+		self.model.add(tt_layer)
+		self.model.add(Reshape((self.nx1, self.nx2)))
+
+	def call(self, inputs):
+		return self.model(inputs)

stnn/pde/circle.py

@@ -0,0 +1,156 @@
+from .common import *
+
+
+def u_dot_thetavec(r, theta, w):
+	"""
+	Dot product of u = (cos(w), sin(w)) with the coordinate vector for theta.
+
+	Args:
+	 	r (float or array-like): the radial coordinate(s)
+		theta (float or array-like): The angular coordinate(s).
+		w (float or array-like): The w coordinate(s).
+
+	Returns:
+		numpy.ndarray: The calculated dot product for each point.
+	"""
+	return r * np.sin(w - theta)
+
+
+def u_dot_thetahat(theta, w):
+	"""
+	Dot product of u = (cos(w), sin(w)) with the unit vector for theta.
+
+	Args:
+		theta (float or array-like): The angular (theta) coordinate(s).
+		w (float or array-like): The w coordinate(s).
+
+	Returns:
+		numpy.ndarray: The calculated dot product for each point.
+	"""
+	return np.sin(w - theta)
+
+
+def u_dot_rvec(theta, w):
+	"""
+	Dot product of u = (cos(w), sin(w)) with the coordinate vector for r.
+
+	Args:
+		theta (float or array-like): The radial (r) coordinate(s).
+		w (float or array-like): The w coordinate(s).
+
+	Returns:
+		numpy.ndarray: The calculated dot product for each point.
+	"""
+	return np.cos(w - theta)
+
+
+def get_system_circle(config):
+	"""
+	For a circular geometry, constructs the matrices, grids, and other quantities corresponding to the PDE system
+	specified by "config"".
+
+	Args:
+		config (dict): Configuration dictionary containing the system parameters.
+
+	Returns:
+		tuple: A tuple containing matrices, grids, etc. for the PDE system
+	"""
+	required_keys = ['nx1', 'nx2', 'nx3', 'ell', 'a2']
+	optional_keys = []
+
+	missing_keys = [key for key in required_keys if key not in config]
+	if missing_keys:
+		raise KeyError(f"Missing keys in config: {', '.join(missing_keys)}")
+
+	unused_keys = [key for key in config if key not in required_keys + optional_keys]
+	if unused_keys:
+		warnings.warn(f"Unused keys in config: {', '.join(unused_keys)}")
+
+	for key in ['nx1', 'nx2', 'nx3']:
+		if not isinstance(config[key], int):
+			raise TypeError(f"{key} must be an integer.")
+
+	for key in ['nx1', 'nx2', 'nx3', 'ell', 'a2']:
+		if config[key] <= 0:
+			raise ValueError(f"{key} must be positive.")
+
+	if config['a2'] < 1.0:
+		raise ValueError('a2 must be greater than 1.')
+
+	nr, ntheta, nw = config['nx1'], config['nx2'], config['nx3']
+	R1 = 1.0
+	R2 = config['a2']
+	ell = config['ell']
+
+	# 1D grids
+	theta, w = get_angular_grids(ntheta, nw)
+	# r grid: non-uniform spacing and Dirichlet boundary conditions
+	y = np.linspace(-np.pi / 2, np.pi / 2, nr + 2)
+	r_ = (R2 - R1) * (np.sin(y) / 2 + 0.5) + R1
+	dr1 = r_[1] - r_[0]
+	dr2 = r_[-1] - r_[-2]
+	r = r_[1:-1]
+
+	# 1D finite-difference operators
+	Dtheta_minus, Dtheta_plus = d_dx_upwind(theta, ntheta)
+	D2w = d2_dx2_fourth_order(w, nw)
+	Dr_minus, Dr_plus = d_dx_upwind_nonuniform(r_, nr)
+
+	# 3D quantities. Kronecker products are used to build the 3D difference operators
+	r_3D, theta_3D, w_3D = np.meshgrid(r, theta, w, indexing = 'ij')
+	I_r = sp.eye(nr)
+	I_theta = sp.eye(ntheta)
+	I_w = sp.eye(nw)
+	Dtheta_3D_minus = sp.kron(sp.kron(I_r, Dtheta_minus), I_w)
+	Dtheta_3D_plus = sp.kron(sp.kron(I_r, Dtheta_plus), I_w)
+	D2w_3D = sp.kron(sp.kron(I_r, I_theta), D2w)
+	Dr_3D_minus = sp.kron(sp.kron(Dr_minus, I_theta), I_w)
+	Dr_3D_plus = sp.kron(sp.kron(Dr_plus, I_theta), I_w)
+
+	# Metric tensor. Note that g_12 = g_21 = 0.
+	g_11 = np.ones_like(r_3D)
+	g_22_over_r = r_3D  # divide out factor of r
+
+	# Dot products
+	dp_r = u_dot_rvec(theta_3D, w_3D)
+	dp_thetahat = u_dot_thetahat(theta_3D, w_3D)
+
+	# Coefficient of d / dr
+	Dr_3D_coeff_meshgrid = dp_r / g_11
+	test_ill_conditioned(Dr_3D_coeff_meshgrid)
+	Dr_3D_coeff = sp.diags(Dr_3D_coeff_meshgrid.ravel())
+
+	# Coefficient of d / dtheta
+	Dtheta_3D_coeff_meshgrid = dp_thetahat / g_22_over_r
+	Dtheta_3D_coeff = sp.diags(Dtheta_3D_coeff_meshgrid.ravel())
+
+	# Upwind differencing
+	Dr_3D_upwind = upwind_operator(Dr_3D_minus, Dr_3D_plus, Dr_3D_coeff_meshgrid)
+	Dtheta_3D_upwind = upwind_operator(Dtheta_3D_minus, Dtheta_3D_plus, Dtheta_3D_coeff_meshgrid)
+
+	# Full operator
+	L = Dr_3D_coeff @ Dr_3D_upwind + Dtheta_3D_coeff @ Dtheta_3D_upwind - (1 / ell) * D2w_3D
+
+	return L, r_3D, theta_3D, w_3D, dr1, dr2, Dr_3D_coeff_meshgrid
+
+
+def get_boundary_quantities_circle(theta_3D, w_3D):
+	"""
+	Gets grid coordinates on the boundaries, as well as slice arrays
+	for positive/negative angles with respect to the boundary angle.
+
+	Args:
+		theta_3D (numpy.ndarray): 3D array of theta values on the grid.
+		w_3D (numpy.ndarray): 3D array of w values on the grid.
+
+	Returns:
+		tuple: Tuple of the grid coordinates and slice arrays
+	"""
+	th1 = theta_3D[0, :, :]
+	wb1 = w_3D[0, :, :]
+	th2 = theta_3D[-1, :, :]
+	wb2 = w_3D[-1, :, :]
+	ib_slice = np.cos(th1 - wb1) > 0
+	ob_slice = np.cos(th2 - wb2) < 0
+
+	return th1, th2, wb1, wb2, ib_slice, ob_slice

stnn/pde/common.py

@@ -0,0 +1,91 @@
+import warnings
+import numpy as np
+import scipy.sparse as sp
+
+
+def d_dx_upwind(x, nx):
+	"""
+	Sparse matrix representation of d/dx using first-order left/right differences with Dirichlet boundary conditions
+	"""
+	dx = x[1] - x[0]
+	Dx_minus = sp.diags([-1, 1], [0, 1], shape = (nx, nx)).tolil() / dx
+	Dx_plus = sp.diags([-1, 1], [-1, 0], shape = (nx, nx)).tolil() / dx
+	Dx_minus[-1, 0] = 1 / dx
+	Dx_plus[0, -1] = -1 / dx
+	Dx_minus = Dx_minus.tocsr()
+	Dx_plus = Dx_plus.tocsr()
+	return Dx_minus, Dx_plus
+
+
+def d2_dx2_fourth_order(x, nx):
+	"""
+	Sparse matrix representation of d^2/dx^2 using fourth order central differences and periodic boundary conditions
+	"""
+	dx = x[1] - x[0]
+	D2x = sp.diags([-1, 16, -30, 16, -1], [-2, -1, 0, 1, 2],
+				   shape = (nx, nx)).tolil() / (12 * dx**2)
+	D2x[0, -1] = 16 / (12 * dx**2)
+	D2x[0, -2] = -1 / (12 * dx**2)
+	D2x[1, -1] = -1 / (12 * dx**2)
+	D2x[-1, 0] = 16 / (12 * dx**2)
+	D2x[-1, 1] = -1 / (12 * dx**2)
+	D2x[-2, 0] = -1 / (12 * dx**2)
+	D2x = D2x.tocsr()
+	return D2x
+
+
+def d_dx_upwind_nonuniform(x, nx):
+	"""
+	Sparse matrix representation of d/dx on a nonuniform grid, using first-order left/right differences 
+	with Dirichlet boundary conditions.
+	"""
+	Dx_ = np.diff(x)
+	Dx_minus = np.diff(x[1:])
+	Dx_minus_inv = 1 / Dx_minus
+	Dx_plus_inv = 1 / Dx_
+	Dx_minus = sp.diags([-Dx_minus_inv, Dx_minus_inv], [0, 1], shape = (nx, nx)).tolil()
+	Dx_plus = sp.diags([-Dx_plus_inv[1:], Dx_plus_inv[:-1]], [-1, 0], shape = (nx, nx)).tolil()
+	Dx_minus = Dx_minus.tocsr()
+	Dx_plus = Dx_plus.tocsr()
+	return Dx_minus, Dx_plus
+
+
+def get_angular_grids(nx2, nx3):
+	"""
+	x2 / x3 grids: uniform spacing and periodic boundary conditions
+	The x3 grid has an offset to ensure cos(x2 - x3) != 0.
+	"""
+	x2 = np.linspace(-np.pi, np.pi, nx2, endpoint = False)
+	x3_min, x3_max = 0 + 0.125 * (2 * np.pi / nx3), 2 * np.pi + 0.125 * (2 * np.pi / nx3)
+	x3 = np.linspace(x3_min, x3_max, nx3, endpoint = False)
+	return x2, x3
+
+
+def upwind_operator(Dx_minus, Dx_plus, Dx_coeff):
+	"""
+	Upwind finite difference operator.
+
+	Args:
+		Dx_minus (scipy.sparse matrix): backward (minus) finite difference operator.
+		Dx_plus (scipy.sparse matrix): forward  (plus) finite difference operator.
+		Dx_coeff (numpy.ndarray): coefficient array
+
+	Returns:
+		scipy.sparse matrix: The upwind operator
+	"""
+	mask_x = Dx_coeff <= 0
+	Dx_masked_minus = sp.diags(mask_x.ravel().astype(int)) @ Dx_minus
+	Dx_masked_plus = sp.diags((~mask_x).ravel().astype(int)) @ Dx_plus
+	Dx_upwind = Dx_masked_minus + Dx_masked_plus
+	return Dx_upwind
+
+
+def test_ill_conditioned(Dx_coeff):
+	"""
+	Test for ill-conditioning. The thresholds are heuristic only.
+	"""
+	ill_conditioning_test = np.min(np.abs(Dx_coeff.ravel()))
+	if ill_conditioning_test < 1e-10:
+		raise ValueError(f'System is ill-conditioned; min |Dx1_coeff| = {ill_conditioning_test}')
+	elif ill_conditioning_test < 1e-6:
+		warnings.warn(f'System may be ill-conditioned; min |Dx1_coeff| = {ill_conditioning_test}')

stnn/pde/ellipse.py

@@ -0,0 +1,154 @@
+from .common import *
+
+
+def u_dot_muvec(mu, eta, w):
+	"""
+	Dot product of u = (cos(w), sin(w)) with the coordinate vector for mu.
+
+	Args:
+		mu (float or array-like): The mu coordinate(s).
+		eta (float or array-like): The eta coordinate(s).
+		w (float or array-like): The w coordinate(s).
+
+	Returns:
+		numpy.ndarray: The calculated dot product for each point.
+	"""
+	return (0.5 * np.cosh(mu) * np.cos(eta - w) + 0.5 * np.cosh(mu) * np.cos(eta + w)
+			+ 0.5 * np.sinh(mu) * np.cos(eta - w) - 0.5 * np.sinh(mu) * np.cos(eta + w))
+
+
+def u_dot_etavec(mu, eta, w):
+	"""
+	Dot product of u = (cos(w), sin(w)) with the coordinate vector for eta.
+
+	Args:
+		mu (float or array-like): The mu coordinate(s).
+		eta (float or array-like): The eta coordinate(s).
+		w (float or array-like): The w coordinate(s).
+
+	Returns:
+		numpy.ndarray: The calculated dot product for each point.
+	"""
+	return (-0.5 * np.sinh(mu) * np.sin(eta + w) - 0.5 * np.sinh(mu) * np.sin(eta - w)
+			+ 0.5 * np.cosh(mu) * np.sin(eta + w) - 0.5 * np.cosh(mu) * np.sin(eta - w))
+
+
+def get_system_ellipse(config):
+	"""
+	For an elliptical geometry, constructs the matrices, grids, and other quantities corresponding to the PDE system
+	specified by "config"".
+
+	Args:
+		config (dict): Configuration dictionary containing the system parameters.
+
+	Returns:
+		tuple: A tuple containing matrices, grids, etc. for the PDE system
+	"""
+	required_keys = ['nx1', 'nx2', 'nx3', 'ell', 'a2', 'eccentricity']
+	optional_keys = []
+
+	missing_keys = [key for key in required_keys if key not in config]
+	if missing_keys:
+		raise KeyError(f"Missing keys in config: {', '.join(missing_keys)}")
+
+	unused_keys = [key for key in config if key not in required_keys + optional_keys]
+	if unused_keys:
+		warnings.warn(f"Unused keys in config: {', '.join(unused_keys)}")
+
+	for key in ['nx1', 'nx2', 'nx3']:
+		if not isinstance(config[key], int):
+			raise TypeError(f"{key} must be an integer.")
+
+	for key in ['nx1', 'nx2', 'nx3', 'ell']:
+		if config[key] <= 0:
+			raise ValueError(f"{key} must be positive.")
+
+	if not (0 <= config['eccentricity'] < 1.0):
+		raise ValueError('eccentricity must be >= 0 and < 1.')
+
+	if config['a2'] <= config['eccentricity']:
+		raise ValueError(f'a2 must be greater than the eccentricity.')
+
+	nmu, neta, nw = config['nx1'], config['nx2'], config['nx3']
+	minor_axis_outer = config['a2']
+	ell = config['ell']
+	minor_axis = 1.0 - config['eccentricity']
+	major_axis = 1.0
+	focal_distance = np.sqrt(major_axis**2 - minor_axis**2)
+	mu1 = np.arccosh(major_axis / focal_distance)
+	major_axis_outer = np.sqrt(focal_distance**2 + minor_axis_outer**2)
+	mu2 = np.arccosh(major_axis_outer / focal_distance)
+
+	# 1D grids
+	eta, w = get_angular_grids(neta, nw)
+	# mu grid: non-uniform spacing and Dirichlet boundary conditions
+	y = np.linspace(-np.pi / 2, np.pi / 2, nmu + 2, dtype = np.float64)
+	mu_ = np.log((np.exp(mu2) - np.exp(mu1)) * (np.sin(y) / 2 + 0.5) + np.exp(mu1))
+	dmu1 = mu_[1] - mu_[0]
+	dmu2 = mu_[-1] - mu_[-2]
+	mu = mu_[1:-1]
+
+	# 1D finite-difference operators
+	Deta_minus, Deta_plus = d_dx_upwind(eta, neta)
+	D2w = d2_dx2_fourth_order(w, nw)
+	Dmu_minus, Dmu_plus = d_dx_upwind_nonuniform(mu_, nmu)
+
+	# 3D quantities. Kronecker products are used to build the 3D difference operators
+	mu_3D, eta_3D, w_3D = np.meshgrid(mu, eta, w, indexing = 'ij')
+	I_mu = sp.eye(nmu)
+	I_eta = sp.eye(neta)
+	I_w = sp.eye(nw)
+	Deta_3D_minus = sp.kron(sp.kron(I_mu, Deta_minus), I_w)
+	Deta_3D_plus = sp.kron(sp.kron(I_mu, Deta_plus), I_w)
+	D2w_3D = sp.kron(sp.kron(I_mu, I_eta), D2w)
+	Dmu_3D_minus = sp.kron(sp.kron(Dmu_minus, I_eta), I_w)
+	Dmu_3D_plus = sp.kron(sp.kron(Dmu_plus, I_eta), I_w)
+
+	# Metric tensor. Note that g_12 = g_21 = 0 and g_11 = g_22.
+	g_11 = focal_distance * (np.cosh(mu_3D) * np.cosh(mu_3D) * np.cos(eta_3D) * np.cos(eta_3D)
+			+ np.sinh(mu_3D) * np.sinh(mu_3D) * np.sin(eta_3D) * np.sin(eta_3D))
+
+	# Dot products
+	dp_mu = u_dot_muvec(mu_3D, eta_3D, w_3D)
+	dp_eta = u_dot_etavec(mu_3D, eta_3D, w_3D)
+
+	# Coefficient of d / dmu
+	Dmu_3D_coeff_meshgrid = dp_mu / g_11
+	test_ill_conditioned(Dmu_3D_coeff_meshgrid)
+
+	Dmu_3D_coeff = sp.diags(Dmu_3D_coeff_meshgrid.ravel())
+
+	# Coefficient of d / deta
+	Deta_3D_coeff_meshgrid = dp_eta / g_11
+	Deta_3D_coeff = sp.diags(Deta_3D_coeff_meshgrid.ravel())
+
+	# Upwind differencing
+	Dmu_3D_upwind = upwind_operator(Dmu_3D_minus, Dmu_3D_plus, Dmu_3D_coeff_meshgrid)
+	Deta_3D_upwind = upwind_operator(Deta_3D_minus, Deta_3D_plus, Deta_3D_coeff_meshgrid)
+
+	# Full operator
+	L = Dmu_3D_coeff @ Dmu_3D_upwind + Deta_3D_coeff @ Deta_3D_upwind - (1 / ell) * D2w_3D
+
+	return L, mu_3D, eta_3D, w_3D, dmu1, dmu2, Dmu_3D_coeff_meshgrid, major_axis_outer
+
+
+def get_boundary_quantities_ellipse(mu_3D, eta_3D, w_3D):
+	"""
+	Gets grid coordinates on the boundaries, as well as slice arrays
+	for positive/negative angles with respect to the boundary angle.
+
+	Args:
+		mu_3D: 3D array of mu values on the grid.
+		eta_3D: 3D array of eta values on the grid.
+		w_3D (numpy.ndarray): 3D array of w values on the grid.
+
+	Returns:
+		tuple: Tuple of the grid coordinates and slice arrays
+	"""
+	eta_2D_ib = eta_3D[0, ...]
+	eta_2D_ob = eta_3D[-1, ...]
+	w_2D_ib = w_3D[0, ...]
+	w_2D_ob = w_3D[-1, ...]
+	ib_slice = u_dot_muvec(mu_3D, eta_3D, w_3D)[0, ...] > 0
+	ob_slice = u_dot_muvec(mu_3D, eta_3D, w_3D)[-1, ...] < 0
+	return eta_2D_ib, eta_2D_ob, w_2D_ib, w_2D_ob, ib_slice, ob_slice

stnn/pde/pde_system.py

@@ -0,0 +1,223 @@
+import numpy as np
+
+from stnn.data.function_generators import generate_random_functions
+from .circle import get_system_circle, get_boundary_quantities_circle, u_dot_thetahat
+from .ellipse import (get_system_ellipse, get_boundary_quantities_ellipse, u_dot_etavec)
+
+
+class PDESystem:
+    """
+    Constructs the PDE system given input parameters. The finite-difference matrices, grids, and other relevant
+    quantities are available as attributes.
+
+    Constructor Args:
+        params (dict): Configuration dictionary containing the parameters that define the PDE system.
+
+    Attributes:
+
+        ib_slice (numpy.ndarray): Boolean array defining nodes adjacent to the inner boundary
+        ob_slice (numpy.ndarray): Boolean array defining nodes adjacent to the outer boundary
+
+        x2_ib (numpy.ndarray): The x2 coordinate values at the inner boundary.
+        x2_ob (numpy.ndarray): The x2 coordinate values at the outer boundary.
+        x3_ib (numpy.ndarray): The x3 coordinate values at the inner boundary.
+        x3_ob (numpy.ndarray): The x3 coordinate values at the outer boundary.
+
+        Dx1_coeff (numpy.ndarray): Coefficients for the advection operator in the radial direction. Used for converting
+                                   boundary conditions to the r.h.s. of the linear system defining a
+                                   boundary-value problem.
+
+        dx1a (numpy.ndarray): Grid spacing adjacent to the inner boundary
+        dx1b (numpy.ndarray): Grid spacing adjacent to the outer boundary
+
+        L (numpy.ndarray): Finite-difference representation of the linear operator defining the PDE.
+
+        x1 (numpy.ndarray): The grid values in radial coordinate (r or mu)
+        x2 (numpy.ndarray): The grid values in the angular coordinate (theta or eta).
+        x3 (numpy.ndarray): The grid values in the w coordinate
+
+        a1 (float): Minor axis of the inner boundary
+        a2 (float): Minor axis of the outer boundary
+        b1 (float): Major axis of the inner boundary
+        b2 (float): Major axis of the outer boundary
+
+        _coords (str): The type of coordinate system used ('ellipse' or 'circle'). This affects how the grids and
+                      other geometric properties are calculated.
+
+        params (dict): The configuration dictionary containing the PDE parameters.
+    """
+    def __init__(self, params):
+        self._required_keys = ['nx1', 'nx2', 'nx3', 'ell', 'a2', 'eccentricity']
+        self._optional_keys = []
+        self.ib_slice = None
+        self.ob_slice = None
+        self.x2_ib = None
+        self.x2_ob = None
+        self.x3_ib = None
+        self.x3_ob = None
+        self.Dx1_coeff = None
+        self.dx1b = None
+        self.dx1a = None
+        self.L = None
+        self.x1 = None
+        self.x2 = None
+        self.x3 = None
+        self.a1 = None
+        self.a2 = None
+        self.b1 = None
+        self.b2 = None
+        self._coords = None
+        self.params = params
+        self.initialize()
+
+    def initialize(self):
+        """
+        Constructs the system matrices and vectors based on the stored configuration.
+
+        Depending on the 'eccentricity' parameter in `self.params`, the coordinate system is set
+        to either 'circle' or 'ellipse'; the domain parametrization and finite-difference grid are defined
+        accordingly.
+        """
+        params = self.params
+        missing_keys = [key for key in self._required_keys if key not in params]
+        if missing_keys:
+            raise KeyError(f"Missing keys in config: {', '.join(missing_keys)}")
+
+        # The functions 'get_system_circle' and 'get_system_ellipse' have a fair amount
+        # of overlap and probably should be combined, but for now they are kept separate
+        # for simplicity and readability.
+        if params['eccentricity'] < 1e-7:
+            self._coords = 'circle'
+            L, x1, x2, x3, dx1a, dx1b, Dx1_coeff = get_system_circle(params)
+            x2_ib, x2_ob, x3_ib, x3_ob, ib_slice, ob_slice = get_boundary_quantities_circle(x2, x3)
+            self.b2 = params['a2']
+        else:
+            self._coords = 'ellipse'
+            (L, x1, x2, x3, dx1a, dx1b, Dx1_coeff, major_axis_outer) = get_system_ellipse(params)
+            x2_ib, x2_ob, x3_ib, x3_ob, ib_slice, ob_slice = get_boundary_quantities_ellipse(x1, x2, x3)
+            self.b2 = major_axis_outer
+
+        self.a1 = 1.0 - params['eccentricity']
+        self.a2 = params['a2']
+        self.b1 = 1.0
+
+        self.x1 = x1
+        self.x2 = x2
+        self.x3 = x3
+
+        self.L = L
+
+        self.dx1a = dx1a
+        self.dx1b = dx1b
+        self.Dx1_coeff = Dx1_coeff
+
+        self.x2_ib = x2_ib
+        self.x2_ob = x2_ob
+        self.x3_ib = x3_ib
+        self.x3_ob = x3_ob
+        self.ib_slice = ib_slice
+        self.ob_slice = ob_slice
+
+    def generate_random_bc(self, func_gen_id):
+        """
+        Generates random boundary conditions for the PDE system.
+
+        Args:
+            func_gen_id (int): Integer representing the type of 'function generator' used to construct the
+            boundary conditions.
+
+        Returns:
+            tuple: A tuple containing:
+                - ibf_data: Inner boundary data
+                - obf_data: Outer boundary data
+                - b: 3D array for passing to the GMRES solver. 'b' contains the boundary data but is defined
+                     on the full 3D grid.
+                - bf: Flattened boundary data before it is permuted
+
+        The boundary conditions are defined on the inner and outer boundaries of the domain and are denoted
+        by 'ibf_data' and 'obf_data'. The function passes 'ibf_data' and 'obf_data' through 'convert_boundary_data',
+        which converts them into formats suitable for passing into the GMRES solver and STNN model (e.g., by
+        reshaping and permutation operations).
+        """
+
+        # Note the change of variable (x2, x3) -> (x2, x2 - x3).
+        ibf_data = generate_random_functions(1, self.x2_ib, self.x2_ib - self.x3_ib,
+											 max_freq=self.params['nx3'], func_gen_id = func_gen_id)[0, self.ib_slice]
+        obf_data = generate_random_functions(1, self.x2_ob, self.x2_ob - self.x3_ob,
+											 max_freq=self.params['nx3'], func_gen_id = func_gen_id)[0, self.ob_slice]
+
+        # Combine boundary data in single vector
+        bf = np.concatenate([ibf_data, obf_data], axis = -1).flatten()
+
+        # Permutes 'ibf_data' and 'obf_data' and construct 'b'
+        ibf_data, obf_data, b = self.convert_boundary_data(ibf_data, obf_data)
+
+        return ibf_data, obf_data, b, bf
+
+    def convert_boundary_data(self, ibf_data, obf_data):
+        """
+        Converts boundary data into formats suitable for passing into the GMRES solver and STNN model.
+
+        Args:
+            ibf_data: Inner boundary data
+            obf_data: Outer boundary data
+
+        Returns:
+            tuple: A tuple containing:
+                - ibf_data: Inner boundary data, permuted to match the input structure of the EinsumTTL layer.
+                - obf_data: Outer boundary data, permuted to match the input structure of the EinsumTTL layer
+                - b: 3D array for passing to the GMRES solver. 'b' contains the boundary data but is defined
+                     on the full 3D grid.
+        """
+        nx1, nx2, nx3 = self.params['nx1'], self.params['nx2'], self.params['nx3']
+        b = np.zeros((nx1, nx2, nx3), dtype=np.float64)
+        b[0, self.ib_slice] = self.Dx1_coeff[0, self.ib_slice] * (ibf_data / self.dx1a)
+        b[-1, self.ob_slice] = -self.Dx1_coeff[-1, self.ob_slice] * (obf_data / self.dx1b)
+
+        if self._coords == 'ellipse':
+            sin_angle = u_dot_etavec(self.x1, self.x2, self.x3)
+        elif self._coords == 'circle':
+            sin_angle = u_dot_thetahat(self.x2, self.x3)
+        else:
+            raise ValueError(f'"_coords" attribute should be either "ellipse" or "circle"; instead received {self._coords}')
+
+        # reshape and permute 'ibf_data' and 'obf_data'
+        sin_angle_i = sin_angle[0, self.ib_slice].reshape(nx2, nx3 // 2)
+        sin_angle_o = sin_angle[-1, self.ob_slice].reshape(nx2, nx3 // 2)
+        W_I = np.argsort(sin_angle_i, axis=1)
+        W_O = np.argsort(sin_angle_o, axis=1)
+        ibf_data = ibf_data.reshape((nx2, nx3 // 2))
+        obf_data = obf_data.reshape((nx2, nx3 // 2))
+        for n in range(nx2):
+            ibf_data[n, :] = ibf_data[n, W_I[n, :]]
+            obf_data[n, :] = obf_data[n, W_O[n, :]]
+
+        return ibf_data, obf_data, b
+
+    def get_xy_grids(self):
+        """
+        Converts the native grids of the PDE system to xy coordinates (no interpolation).
+
+        The function also applies a wrap-around in the x2 domain for plotting purposes, ensuring
+        continuity across the (periodic) domain.
+
+        Returns:
+            tuple of numpy.ndarray: A tuple containing two 2D numpy arrays:
+                - x_grid: The x-coordinates grid.
+                - y_grid: The y-coordinates grid.
+        """
+        x1_1D = self.x1[:, 0, 0]
+        x2_1D = self.x2[0, :, 0]
+        x2_1D = np.append(x2_1D, np.array([np.pi - 1e-3]))  # wrap around for plotting
+        x1_2D, x2_2D = np.meshgrid(x1_1D, x2_1D, indexing='ij')
+        if self._coords == 'ellipse':
+            focal_distance = np.sqrt(self.b1**2 - self.a1**2)
+            x_grid = focal_distance * np.sinh(x1_2D) * np.cos(x2_2D)
+            y_grid = focal_distance * np.cosh(x1_2D) * np.sin(x2_2D)
+        elif self._coords == 'circle':
+            x_grid = x1_2D * np.cos(x2_2D)
+            y_grid = x1_2D * np.sin(x2_2D)
+        else:
+            raise ValueError(f'"_coords" should be either "ellipse" or "circle"; instead received {self._coords}')
+
+        return x_grid, y_grid

stnn/tests/test_circle.py

@@ -0,0 +1,63 @@
+import unittest
+import copy
+import numpy as np
+import scipy.sparse as sp
+from stnn.pde.circle import get_system_circle
+
+
+class TestGetSystemCircle(unittest.TestCase):
+
+	def setUp(self):
+		self.config = {
+			'nx1': 10,
+			'nx2': 20,
+			'nx3': 30,
+			'a2': 2.0,
+			'ell': 1.5,
+		}
+		self.saved_config = copy.deepcopy(self.config)
+		self._required_keys = ['nx1', 'nx2', 'nx3', 'ell', 'a2']
+		self._optional_keys = []
+
+	def test_valid_output(self):
+		L, r_3D, theta_3D, w_3D, dr1, dr2, Dr_3D_coeff_meshgrid = get_system_circle(self.config)
+
+		# Test types
+		self.assertIsInstance(L, sp.csr_matrix)
+		self.assertIsInstance(r_3D, np.ndarray)
+		self.assertIsInstance(theta_3D, np.ndarray)
+		self.assertIsInstance(w_3D, np.ndarray)
+		self.assertIsInstance(Dr_3D_coeff_meshgrid, np.ndarray)
+
+		# Test shapes
+		self.assertEqual(r_3D.shape, (self.config['nx1'], self.config['nx2'], self.config['nx3']))
+		self.assertEqual(theta_3D.shape, (self.config['nx1'], self.config['nx2'], self.config['nx3']))
+		self.assertEqual(w_3D.shape, (self.config['nx1'], self.config['nx2'], self.config['nx3']))
+
+		# Test values
+		self.assertTrue(dr1 > 0)
+		self.assertTrue(dr2 > 0)
+
+	def test_missing_keys(self):
+		for key in self._required_keys:
+			del self.config[key]
+			with self.assertRaises(KeyError):
+				get_system_circle(self.config)
+			self.config[key] = self.saved_config[key]
+
+	def test_invalid_parameters(self):
+		for key in self._required_keys:
+			self.config[key] = 0  # None of the required keys should be zero.
+			with self.assertRaises(ValueError):
+				get_system_circle(self.config)
+			self.config[key] = self.saved_config[key]
+
+	def test_unused_params_warning(self):
+		copied_params = copy.deepcopy(self.config)
+		copied_params['unusedkey'] = 0
+		with self.assertWarns(UserWarning) as _:
+			get_system_circle(copied_params)
+
+
+if __name__ == '__main__':
+	unittest.main()

stnn/tests/test_dependencies.py

@@ -0,0 +1,25 @@
+import unittest
+import importlib
+
+
+class TestDependencies(unittest.TestCase):
+	def test_required_modules_installed(self):
+		required_modules = [
+			'numpy',
+			'tensorflow',
+			'numpy',
+			't3f',
+			'scipy',
+			'h5py',
+			'matplotlib',
+			'pydot',
+			'openvino',
+		]
+
+		for module in required_modules:
+			with self.subTest(module = module):
+				importlib.import_module(module)
+
+
+if __name__ == '__main__':
+	unittest.main()

stnn/tests/test_differential_ops.py

@@ -0,0 +1,115 @@
+import unittest
+import numpy as np
+from stnn.pde.common import d_dx_upwind, d2_dx2_fourth_order, d_dx_upwind_nonuniform
+
+
+def assert_derivative(operator, function, expected_derivative, boundary = None, rtol = None, atol = None):
+	"""
+	Assert the derivative of a function using a given finite-difference operator
+
+	Args:
+		operator (np.ndarray or sparse matrix): Differential operator matrix.
+		function (np.ndarray): Values of the function to differentiate
+		expected_derivative (np.ndarray):: Expected result of the derivative.
+		boundary (int): Specifies if boundary elements should be excluded, and how many.
+		rtol: Relative tolerance.
+		atol: Absolute tolerance.
+	"""
+	observed_derivative = operator @ function
+	if boundary is not None:
+		observed_derivative = observed_derivative[boundary:-boundary]
+		expected_derivative = expected_derivative[boundary:-boundary]
+
+	np.testing.assert_allclose(observed_derivative, expected_derivative, rtol = rtol, atol = atol)
+
+
+class TestDifferentialOperators(unittest.TestCase):
+
+	def setUp(self):
+		# ----- Set up grids
+
+		# grid for radial coordinate (non-periodic, non-uniform spacing)
+		self.nx = 100000
+		zs = np.linspace(-np.pi / 2, np.pi / 2, self.nx + 2)
+		R1, R2 = 0.5, 1.4
+		self.z_ = (R2 - R1) * (np.sin(zs) / 2 + 0.5) + R1
+		self.z = self.z_[1:-1]
+
+		# grid for angular coordinates (periodic, uniform spacing)
+		self.y = np.linspace(0, 2 * np.pi, self.nx, endpoint = False)
+		dy = (2 * np.pi) / self.nx
+
+		# ----- Set tolerances
+
+		# Tolerances for "exact" tests, i.e., where the finite differences do not have truncation error
+		self.atol = 1e-6
+		self.rtol = 1e-6
+
+		# Tolerances for some inexact tests
+		self.atol_inexact = 1e-3
+		self.rtol_z = 3 * np.max(np.diff(self.z_))  # relative tolerance of 3*dx for first-order one-sided differences
+		self.rtol_y = 3 * dy  # relative tolerance of 3*dx for first-order one-sided differences
+		self.rtol_y2 = 3 * dy**4  # relative tolerance of 3*dy**4 for fourth-order central differences
+
+		# ----- Test inputs
+
+		# "Exact" test inputs
+		self.f1 = -2.3 * np.ones(self.nx)
+		self.f2 = 0.8 * self.z
+		self.f3 = 0.8 * self.y
+		self.f4 = -0.1 * self.y * self.y
+
+		# Inexact test inputs
+		# noinspection PyRedundantParentheses
+		self.g1 = (self.z)**2 - (self.z)**3
+		self.g2 = (np.cos(self.y))**2 - (np.sin(self.y))**3
+		self.g3 = (np.sin(2 * self.y))**2 - (np.cos(self.y))**3
+
+	def test_d_dx_upwind(self):
+		dx_m, dx_p = d_dx_upwind(self.y, self.nx)
+
+		assert_derivative(dx_m, self.f1, np.zeros(self.nx), rtol = self.rtol, atol = self.atol)
+		assert_derivative(dx_p, self.f1, np.zeros(self.nx), rtol = self.rtol, atol = self.atol)
+
+		assert_derivative(dx_m, self.f3, 0.8 * np.ones(self.nx), boundary = 1, rtol = self.rtol, atol = self.atol)
+		assert_derivative(dx_p, self.f3, 0.8 * np.ones(self.nx), boundary = 1, rtol = self.rtol, atol = self.atol)
+
+		expected_dg2 = -2 * np.cos(self.y) * np.sin(self.y) - 3 * np.sin(self.y)**2 * np.cos(self.y)
+		assert_derivative(dx_m, self.g2, expected_dg2, rtol = self.rtol_y, atol = self.atol_inexact)
+		assert_derivative(dx_p, self.g2, expected_dg2, rtol = self.rtol_y, atol = self.atol_inexact)
+
+		expected_dg3 = 4 * np.sin(2 * self.y) * np.cos(2 * self.y) + 3 * np.cos(self.y)**2 * np.sin(self.y)
+		assert_derivative(dx_m, self.g3, expected_dg3, rtol = self.rtol_y, atol = self.atol_inexact)
+		assert_derivative(dx_p, self.g3, expected_dg3, rtol = self.rtol_y, atol = self.atol_inexact)
+
+	def test_d2_dx2_fourth_order(self):
+		d2x = d2_dx2_fourth_order(self.y, self.nx)
+
+		assert_derivative(d2x, self.f1, np.zeros(self.nx), rtol = self.rtol, atol = self.atol)
+		assert_derivative(d2x, self.f3, np.zeros(self.nx), boundary = 4, rtol = self.rtol, atol = self.atol)
+		assert_derivative(d2x, self.f4, -0.2 * np.ones(self.nx), boundary = 4, rtol = self.rtol, atol = self.atol)
+
+		expected_dg2 = 2 * np.sin(self.y)**2 - 2 * np.cos(self.y)**2 - \
+					   6 * np.sin(self.y) * np.cos(self.y)**2 + 3 * np.sin(self.y)**3
+		assert_derivative(d2x, self.g2, expected_dg2, rtol = self.rtol_y2, atol = self.atol_inexact)
+
+		expected_dg3 = 8 * np.cos(2 * self.y)**2 - 8 * np.sin(2 * self.y)**2 - \
+					   6 * np.cos(self.y) * np.sin(self.y)**2 + 3 * np.cos(self.y)**3
+		assert_derivative(d2x, self.g3, expected_dg3, rtol = self.rtol_y2, atol = self.atol_inexact)
+
+	def test_d_dx_upwind_nonuniform(self):
+		dx_minus, dx_plus = d_dx_upwind_nonuniform(self.z_, self.nx)
+
+		assert_derivative(dx_minus, self.f1, np.zeros(self.nx), boundary = 1, rtol = self.rtol, atol = self.atol)
+		assert_derivative(dx_plus, self.f1, np.zeros(self.nx), boundary = 1, rtol = self.rtol, atol = self.atol)
+
+		assert_derivative(dx_minus, self.f2, 0.8 * np.ones(self.nx), boundary = 1, rtol = self.rtol, atol = self.atol)
+		assert_derivative(dx_plus, self.f2, 0.8 * np.ones(self.nx), boundary = 1, rtol = self.rtol, atol = self.atol)
+
+		expected_dg1 = (2 * self.z - 3 * self.z**2)
+		assert_derivative(dx_minus, self.g1, expected_dg1, boundary = 1, rtol = self.rtol_z, atol = self.atol_inexact)
+		assert_derivative(dx_plus, self.g1, expected_dg1, boundary = 1, rtol = self.rtol_z, atol = self.atol_inexact)
+
+
+if __name__ == '__main__':
+	unittest.main()

stnn/tests/test_ellipse.py

@@ -0,0 +1,67 @@
+import unittest
+import copy
+import numpy as np
+import scipy.sparse as sp
+from stnn.pde.ellipse import get_system_ellipse
+
+
+class TestGetSystemCircle(unittest.TestCase):
+
+	def setUp(self):
+		self.config = {
+			'nx1': 10,
+			'nx2': 20,
+			'nx3': 30,
+			'a2': 2.0,
+			'ell': 1.5,
+			'eccentricity': 0.5
+		}
+		self.saved_config = copy.deepcopy(self.config)
+		self._required_keys = ['nx1', 'nx2', 'nx3', 'ell', 'a2', 'eccentricity']
+		self._optional_keys = []
+
+	def test_valid_output(self):
+		L, mu_3D, eta_3D, w_3D, dmu1, dmu2, Dmu_3D_coeff_meshgrid, major_axis_outer = get_system_ellipse(self.config)
+
+		# Test types
+		self.assertIsInstance(L, sp.csr_matrix)
+		self.assertIsInstance(mu_3D, np.ndarray)
+		self.assertIsInstance(eta_3D, np.ndarray)
+		self.assertIsInstance(w_3D, np.ndarray)
+		self.assertIsInstance(Dmu_3D_coeff_meshgrid, np.ndarray)
+
+		# Test shapes
+		self.assertEqual(mu_3D.shape, (self.config['nx1'], self.config['nx2'], self.config['nx3']))
+		self.assertEqual(eta_3D.shape, (self.config['nx1'], self.config['nx2'], self.config['nx3']))
+		self.assertEqual(w_3D.shape, (self.config['nx1'], self.config['nx2'], self.config['nx3']))
+		N = self.config['nx1'] * self.config['nx2'] * self.config['nx3']
+		self.assertEqual(L.shape, (N, N))
+
+		# Test values
+		self.assertTrue(dmu1 > 0)
+		self.assertTrue(dmu2 > 0)
+		self.assertTrue(major_axis_outer > self.config['a2'])
+
+	def test_missing_keys(self):
+		for key in self._required_keys:
+			del self.config[key]
+			with self.assertRaises(KeyError):
+				get_system_ellipse(self.config)
+			self.config[key] = self.saved_config[key]
+
+	def test_invalid_parameters(self):
+		for key in self._required_keys:
+			self.config[key] = -1  # None of the required keys should be negative.
+			with self.assertRaises(ValueError):
+				get_system_ellipse(self.config)
+			self.config[key] = self.saved_config[key]
+
+	def test_unused_params_warning(self):
+		copied_params = copy.deepcopy(self.config)
+		copied_params['unusedkey'] = 0
+		with self.assertWarns(UserWarning) as _:
+			get_system_ellipse(copied_params)
+
+
+if __name__ == '__main__':
+	unittest.main()

stnn/tests/test_file.py

@@ -0,0 +1,121 @@
+import unittest
+import numpy as np
+import h5py
+import tempfile
+from stnn.data.preprocessing import get_data_from_file, load_data, load_training_data
+
+
+class TestGetDataFromFile(unittest.TestCase):
+
+	def setUp(self):
+		self.temp_file = tempfile.NamedTemporaryFile(delete = False)
+		self.nx1, self.nx2, self.nx3 = 30, 20, 16
+		self.Nsamples = 10
+		with h5py.File(self.temp_file.name, 'w') as f:
+			f.create_dataset('ell', data = np.random.rand(self.Nsamples))
+			f.create_dataset('a1', data = np.random.rand(self.Nsamples))
+			f.create_dataset('a2', data = np.random.rand(self.Nsamples))
+			f.create_dataset('rho', data = np.random.rand(self.Nsamples, self.nx1, self.nx2))
+			f.create_dataset('ibf', data = np.random.rand(self.Nsamples, self.nx2, self.nx3 // 2))
+			f.create_dataset('obf', data = np.random.rand(self.Nsamples, self.nx2, self.nx3 // 2))
+
+		self.temp_file1 = tempfile.NamedTemporaryFile(delete = False)
+		with h5py.File(self.temp_file1.name, 'w') as f:
+			f.create_dataset('ell', data = np.random.rand(self.Nsamples))
+			f.create_dataset('a1', data = np.random.rand(self.Nsamples))
+			f.create_dataset('a2', data = np.random.rand(self.Nsamples))
+			f.create_dataset('rho', data = np.random.rand(self.Nsamples, self.nx1, self.nx2))
+			f.create_dataset('ibf', data = np.random.rand(self.Nsamples, self.nx2, self.nx3 // 2))
+			f.create_dataset('obf', data = np.random.rand(self.Nsamples, self.nx2, self.nx3 // 2))
+
+		self.bad_file = tempfile.NamedTemporaryFile(delete = False)
+		with h5py.File(self.bad_file.name, 'w') as f:
+			f.create_dataset('ell', data = np.random.rand(self.Nsamples))
+			f.create_dataset('a1', data = np.random.rand(self.Nsamples))
+			f.create_dataset('a2', data = np.random.rand(self.Nsamples))
+			f.create_dataset('rho', data = np.random.rand(self.Nsamples, self.nx1, self.nx2))
+
+	def tearDown(self):
+		self.temp_file.close()
+		self.temp_file1.close()
+		self.bad_file.close()
+
+	def test_missing_datasets(self):
+		with self.assertRaises(ValueError):
+			get_data_from_file(self.bad_file.name, self.nx2, self.nx2)
+
+	def test_data_extraction_shapes(self):
+		result = get_data_from_file(self.temp_file.name, self.nx2, self.nx3)
+		self.assertEqual(result[0].shape, (self.Nsamples,))
+		self.assertEqual(result[1].shape, (self.Nsamples,))
+		self.assertEqual(result[2].shape, (self.Nsamples,))
+		self.assertEqual(result[3].shape, (self.Nsamples, 2 * self.nx2, self.nx3 // 2))
+		self.assertEqual(result[4].shape, (self.Nsamples, self.nx1, self.nx2))
+
+	def test_nrange_parameter(self):
+		Nrange = (2, 5)
+		result = get_data_from_file(self.temp_file.name, self.nx2, self.nx3, Nrange = Nrange)
+		expected_size = Nrange[1] - Nrange[0]
+		self.assertEqual(result[0].shape, (expected_size,))
+		self.assertEqual(result[0].shape, (expected_size,))
+
+	def test_list_input(self):
+		file_list = [self.temp_file.name, self.temp_file1.name]
+		Nrange_list = [(0, -1), (0, -1)]
+		with self.assertRaises(TypeError):
+			# noinspection PyTypeChecker
+			_ = get_data_from_file(file_list, self.nx2, self.nx3, Nrange = Nrange_list)
+
+	def test_invalid_Nrange(self):
+		Nrange_list = [(0, -1), (0, -1)]
+		with self.assertRaises(TypeError):
+			_ = get_data_from_file(self.temp_file.name, self.nx2, self.nx3, Nrange = Nrange_list)
+
+		for Nrange in [(0, 1, 1), 1, (1), (1.5, 3), (3, 1.5), (1.5, 1.5), 'x']:
+			with self.assertRaises(TypeError):
+				_ = get_data_from_file(self.temp_file.name, self.nx2, self.nx3, Nrange = Nrange)
+			with self.assertRaises(TypeError):
+				_ = get_data_from_file(self.temp_file.name, self.nx2, self.nx3, Nrange = list(Nrange))
+
+	def test_good_data_load(self):
+		files = [self.temp_file.name, self.temp_file1.name]
+		Nrange_list = [(0, None), (0, self.Nsamples)]
+		ell1, ell2, a1, a2 = 0.1, 2.0, 1.0, 5.0
+
+		params, bf, rho = load_data(files, self.nx2, self.nx3, ell1, ell2, a1, a2, Nrange_list = Nrange_list)
+		self.assertEqual(params.shape, (2 * self.Nsamples, 3))
+		self.assertEqual(bf.shape, (2 * self.Nsamples, 2 * self.nx2, self.nx3 // 2))
+		self.assertEqual(rho.shape, (2 * self.Nsamples, self.nx1, self.nx2))
+
+		test_size = 0.3
+		(params_train, bf_train, rho_train,
+		 params_test, bf_test, rho_test) = load_training_data(files, self.nx2, self.nx3,
+															  ell1, ell2, a1, a2, test_size = test_size,
+															  Nrange_list = Nrange_list)
+		Ntest = int(test_size * 2 * self.Nsamples)
+		Ntrain = 2 * self.Nsamples - Ntest
+		self.assertEqual(params_train.shape, (Ntrain, 3))
+		self.assertEqual(bf_train.shape, (Ntrain, 2 * self.nx2, self.nx3 // 2))
+		self.assertEqual(rho_train.shape, (Ntrain, self.nx1, self.nx2))
+		self.assertEqual(params_test.shape, (Ntest, 3))
+		self.assertEqual(bf_test.shape, (Ntest, 2 * self.nx2, self.nx3 // 2))
+		self.assertEqual(rho_test.shape, (Ntest, self.nx1, self.nx2))
+
+	def test_bad_data_load(self):
+		files = [self.temp_file.name, self.temp_file1.name]
+		Nrange_list = (0, -1)
+		ell1, ell2, a1, a2 = 0.1, 2.0, 1.0, 5.0
+		with self.assertRaises(TypeError):
+			_ = load_data(files, self.nx2, self.nx3, ell1, ell2, a1, a2, Nrange_list = Nrange_list)
+		with self.assertRaises(TypeError):
+			_ = load_data(files, self.nx2, self.nx3, ell1, ell2, a1, a2, Nrange_list = list(Nrange_list))
+
+		Nrange_list = [(0, -1), (0, -1)]
+		for test_size in [-1, 0.0, 1.5]:
+			with self.assertRaises(ValueError):
+				_ = load_training_data(files, self.nx2, self.nx3,
+									   ell1, ell2, a1, a2, test_size = test_size, Nrange_list = Nrange_list)
+
+
+if __name__ == '__main__':
+	unittest.main()

stnn/tests/test_pde_system.py

@@ -0,0 +1,54 @@
+import copy
+import unittest
+
+import numpy as np
+
+from stnn.pde.pde_system import PDESystem
+
+
+class TestPDESystem(unittest.TestCase):
+
+	def setUp(self):
+		self.params = {
+			'nx1': 50,
+			'nx2': 100,
+			'nx3': 75,
+			'a1': 1.0,
+			'a2': 2.0,
+			'ell': 0.1,
+			'eccentricity': 0.5
+		}
+		self.system = PDESystem(self.params)
+
+	def test_initialization(self):
+		self.assertEqual(self.system.params, self.params)
+
+	def test_attribute_types(self):
+		self.assertIsInstance(self.system.ib_slice, np.ndarray)
+		self.assertIsInstance(self.system.ob_slice, np.ndarray)
+
+	def test_attribute_values(self):
+		self.assertEqual(self.system.a1, 1.0 - self.params['eccentricity'])
+
+	def test_coordinate_system(self):
+		expected_coords = 'ellipse' if self.params['eccentricity'] != 0 else 'circle'
+		self.assertEqual(self.system._coords, expected_coords)
+
+	def test_unused_params_warning(self):
+		copied_params = copy.deepcopy(self.params)
+		copied_params['unusedkey'] = 0
+		with self.assertWarns(UserWarning) as _:
+			_ = PDESystem(copied_params)
+
+	def test_L(self):
+		# If f(x1, x2, x3) is constant, then L * f should be 0 except adjacent to the boundary.
+		L = self.system.L
+		nx1, nx2, nx3 = self.params['nx1'], self.params['nx2'], self.params['nx3']
+		f0 = 2.3 * np.ones((nx1, nx2, nx3))
+		result = L @ f0.ravel()
+		result = result.reshape((nx1, nx2, nx3))
+		np.testing.assert_allclose(result[1:-1, ...], 0, atol = 1e-7)
+
+
+if __name__ == '__main__':
+	unittest.main()

stnn/tests/test_preprocessing.py

@@ -0,0 +1,99 @@
+import unittest
+import numpy as np
+from stnn.data.preprocessing import train_test_split
+
+
+class TestTrainTestSplit(unittest.TestCase):
+
+	def setUp(self):
+		self.X_array = np.array([[1, 2], [3, 4], [5, 6], [7, 8]])
+		self.Y_array = np.array([1, 2, 3, 4])
+		self.X_list = [self.X_array, self.X_array]
+		self.Y_list = [self.Y_array, self.Y_array]
+		self.X_list_bad = [self.Y_array, self.X_array]
+		self.Y_list_bad = [self.Y_array, self.X_array]
+
+	def test_basic_functionality_array(self):
+		X_train, X_test, Y_train, Y_test = train_test_split(self.X_array, self.Y_array, test_size = 0.25)
+		self.assertEqual(len(X_train), 3)
+		self.assertEqual(len(X_test), 1)
+		self.assertEqual(len(Y_train), 3)
+		self.assertEqual(len(Y_test), 1)
+
+	def test_basic_functionality_list(self):
+		X_train, X_test, Y_train, Y_test = train_test_split(self.X_list, self.Y_list, test_size = 0.25)
+		self.assertEqual(len(X_train[0]), 3)
+		self.assertEqual(len(X_test[0]), 1)
+		self.assertEqual(len(Y_train[0]), 3)
+		self.assertEqual(len(Y_test[0]), 1)
+
+	def test_return_type_consistency_array(self):
+		X_train, X_test, Y_train, Y_test = train_test_split(self.X_array, self.Y_array, test_size = 0.25)
+		self.assertIsInstance(X_train, np.ndarray)
+		self.assertIsInstance(X_test, np.ndarray)
+		self.assertIsInstance(Y_train, np.ndarray)
+		self.assertIsInstance(Y_test, np.ndarray)
+
+		X_train, X_test, Y_train, Y_test = train_test_split([self.X_array], [self.Y_array], test_size = 0.25)
+		self.assertIsInstance(X_train, list)
+		self.assertIsInstance(X_test, list)
+		self.assertIsInstance(Y_train, list)
+		self.assertIsInstance(Y_test, list)
+
+	def test_return_type_consistency_list(self):
+		X_train, X_test, Y_train, Y_test = train_test_split(self.X_list, self.Y_list, test_size = 0.25)
+		self.assertIsInstance(X_train, list)
+		self.assertIsInstance(X_test, list)
+		self.assertIsInstance(Y_train, list)
+		self.assertIsInstance(Y_test, list)
+
+		# noinspection PyTypeChecker
+		X_train, X_test, Y_train, Y_test = train_test_split(tuple(self.X_list), tuple(self.Y_list), test_size = 0.25)
+		self.assertIsInstance(X_train, list)
+		self.assertIsInstance(X_test, list)
+		self.assertIsInstance(Y_train, list)
+		self.assertIsInstance(Y_test, list)
+
+	def test_random_state(self):
+		X_train1, X_test1, Y_train1, Y_test1 = train_test_split(self.X_array, self.Y_array, test_size = 0.25,
+																random_state = 42)
+		X_train2, X_test2, Y_train2, Y_test2 = train_test_split(self.X_array, self.Y_array, test_size = 0.25,
+																random_state = 42)
+		np.testing.assert_array_equal(X_train1, X_train2)
+		np.testing.assert_array_equal(X_test1, X_test2)
+		np.testing.assert_array_equal(Y_train1, Y_train2)
+		np.testing.assert_array_equal(Y_test1, Y_test2)
+
+	def test_invalid_test_size(self):
+		with self.assertRaises(ValueError):
+			train_test_split(self.X_array, self.Y_array, test_size = -0.1)
+		with self.assertRaises(ValueError):
+			train_test_split(self.X_array, self.Y_array, test_size = 1.5)
+
+	def test_inconsistent_length(self):
+		X = np.array([[1, 2], [3, 4]])
+		Y = np.array([1, 2, 3])
+		with self.assertRaises(ValueError):
+			train_test_split(X, Y)
+		with self.assertRaises(ValueError):
+			train_test_split(self.X_list_bad, self.Y_list_bad)
+		with self.assertRaises(ValueError):
+			train_test_split(self.X_list_bad, self.Y_list)
+		with self.assertRaises(ValueError):
+			train_test_split(self.X_list, self.Y_list_bad)
+
+	def test_empty(self):
+		X_empty = np.zeros(0)
+		Y_empty = np.zeros(0)
+		with self.assertRaises(ValueError):
+			train_test_split(X_empty, Y_empty)
+		with self.assertRaises(ValueError):
+			train_test_split([X_empty], [])
+		with self.assertRaises(ValueError):
+			train_test_split([], [Y_empty])
+		with self.assertRaises(ValueError):
+			train_test_split([X_empty], [Y_empty])
+
+
+if __name__ == '__main__':
+	unittest.main()

stnn/tests/test_stats.py

@@ -0,0 +1,41 @@
+import numpy as np
+import unittest
+import os
+from stnn.utils.stats import get_stats
+
+
+class TestGetStats(unittest.TestCase):
+
+	def setUp(self):
+		self.rho = np.array([[1, 2], [3, 4]])
+		self.rho_pred = np.array([[1, 2], [3, 4]])
+		self.filename = 'test_stats.npz'
+
+	def test_correctness(self):
+		get_stats(self.rho, self.rho_pred, self.filename)
+		with np.load(self.filename) as data:
+			self.assertAlmostEqual(data['max_loss'], 0.0, places=5)
+			self.assertEqual(data['avg_loss'], 0.0)
+			self.assertEqual(data['N'], self.rho.shape[0])
+
+	def test_file_creation(self):
+		get_stats(self.rho, self.rho_pred, self.filename)
+		self.assertTrue(os.path.exists(self.filename))
+
+	def test_file_content(self):
+		get_stats(self.rho, self.rho_pred, self.filename)
+		with np.load(self.filename) as data:
+			self.assertIn('max_loss', data)
+			self.assertIn('avg_loss', data)
+			self.assertIn('N', data)
+
+	def test_invalid_input(self):
+		with self.assertRaises(ValueError):
+			get_stats(np.array([1, 2]), np.array([[1, 2], [3, 4]]))
+
+	def tearDown(self):
+		if os.path.exists(self.filename):
+			os.remove(self.filename)
+
+if __name__ == '__main__':
+	unittest.main()

stnn/tests/test_stnn_config.py

@@ -0,0 +1,61 @@
+import unittest
+import numpy as np
+import copy
+from stnn.nn.stnn import build_stnn
+
+
+class TestBuildSTNN(unittest.TestCase):
+
+	def setUp(self):
+		self.config = {
+			'K': 1,
+			'nx1': 8,
+			'nx2': 8,
+			'nx3': 8,
+			'd': 8,
+			'W': 3,
+			'shape1': [1, 2, 3],
+			'shape2': [2, 2, 2],
+			'ranks': [1, 2, 2, 1],
+		}
+		self.saved_config = copy.deepcopy(self.config)
+		self._required_keys = ['nx1', 'nx2', 'nx3', 'K', 'd', 'shape1','shape2','ranks','W']
+		self._optional_keys = ['use_regularization', 'regularization_strength']
+
+	def test_missing_keys(self):
+		for key in self._required_keys:
+			del self.config[key]
+			with self.assertRaises(KeyError):
+				build_stnn(self.config)
+			self.config[key] = self.saved_config[key]
+
+	def test_invalid_values(self):
+		for key in ['K', 'd', 'W', 'nx1', 'nx2', 'nx3']:
+			for value in [1.5, 'a', None, np.nan]:
+				with self.subTest(value = value):
+					self.config[key] = value
+					with self.assertRaises(TypeError):
+						build_stnn(self.config)
+					self.config[key] = self.saved_config[key]
+			value = -1
+			with self.subTest(value = value):
+				self.config[key] = value
+				with self.assertRaises(ValueError):
+					build_stnn(self.config)
+				self.config[key] = self.saved_config[key]
+
+		self.config['nx3'] = 7  # not divisible by 2
+		with self.assertRaises(ValueError):
+			build_stnn(self.config)
+		self.config[key] = self.saved_config[key]
+
+	def test_positive_values(self):
+		for key in ['K', 'nx1', 'nx2', 'nx3', 'd']:
+			self.config[key] = 0
+			with self.assertRaises(ValueError):
+				build_stnn(self.config)
+			self.config[key] = self.saved_config[key]
+
+
+if __name__ == '__main__':
+	unittest.main()

stnn/tests/test_ttl.py

@@ -0,0 +1,51 @@
+import unittest
+from stnn.nn.stnn_layers import TTL
+
+
+class TestTTL(unittest.TestCase):
+
+	def test_missing_config_keys(self):
+		with self.assertRaises(KeyError):
+			TTL(config = {})  # Empty config
+
+	def test_invalid_use_regularization_type(self):
+		config = {'use_regularization': 'not a boolean', 'nx1': 1, 'nx2': 1, 'nx3': 1, 'shape1': [1], 'shape2': [1],
+				  'ranks': [1], 'W': 1}
+		with self.assertRaises(TypeError):
+			TTL(config = config)
+
+	def test_invalid_nx_values(self):
+		config = {'nx1': 5, 'nx2' : 5, 'nx3': 8, 'use_regularization': False, 'shape1': [1], 'shape2': [1], 'ranks': [1], 'W': 1}
+		for nx in ['nx1', 'nx2', 'nx3']:
+			for value in [-1, 0]:
+				config.update({nx: value})
+				with self.assertRaises(ValueError):
+					TTL(config = config)
+
+	def test_invalid_W_value(self):
+		config = {'use_regularization': False, 'nx1': 1, 'nx2': 1, 'nx3': 1, 'shape1': [1], 'shape2': [1], 'ranks': [1],
+				  'W': -1}
+		with self.assertRaises(ValueError):
+			TTL(config = config)
+
+	def test_shape_length_mismatch(self):
+		config = {'use_regularization': False, 'nx1': 1, 'nx2': 1, 'nx3': 1, 'shape1': [1, 2], 'shape2': [1],
+				  'ranks': [1], 'W': 1}
+		with self.assertRaises(ValueError):
+			TTL(config = config)
+
+	def test_incorrect_shape1_product(self):
+		config = {'use_regularization': False, 'nx1': 2, 'nx2': 2, 'nx3': 1, 'shape1': [1, 3], 'shape2': [4],
+				  'ranks': [1], 'W': 1}
+		with self.assertRaises(ValueError):
+			TTL(config = config)
+
+	def test_incorrect_shape2_product(self):
+		config = {'use_regularization': False, 'nx1': 1, 'nx2': 2, 'nx3': 1, 'shape1': [2], 'shape2': [1, 5],
+				  'ranks': [1], 'W': 1}
+		with self.assertRaises(ValueError):
+			TTL(config = config)
+
+
+if __name__ == '__main__':
+	unittest.main()

stnn/utils/input_output.py

@@ -0,0 +1,135 @@
+import h5py
+from scipy.sparse import csr_matrix
+import numpy as np
+import json
+import tensorflow as tf
+from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
+
+
+def save_to_hdf5(filename, datasets, start_idx, end_idx):
+	"""
+	Saves subsets of datasets to an HDF5 file, either by creating new datasets or appending to existing ones.
+
+	Args:
+		filename (str): The name of the HDF5 file where the data will be saved.
+		datasets (dict): A dictionary where keys are dataset names and values are the corresponding data arrays.
+		start_idx (int): The starting index of the data slice to be saved.
+		end_idx (int): The ending index (exclusive) of the data slice to be saved.
+
+	This function will create new datasets if they do not already exist. If a dataset already exists, it will be resized
+	to accommodate the new data, and the data slice will be appended.
+	"""
+	print(f'Saving data from n = {start_idx} to n = {end_idx}...')
+	with h5py.File(filename, 'a') as f:
+		for name, data in datasets.items():
+			if name not in f:
+				f.create_dataset(name, data = data[start_idx:end_idx], maxshape = (None,) + data.shape[1:],
+								 chunks = (1,) + data.shape[1:])
+			else:
+				f[name].resize((f[name].shape[0] + end_idx - start_idx,) + f[name].shape[1:])
+				f[name][-(end_idx - start_idx):] = data[start_idx:end_idx]
+	print('Done.')
+
+
+def write_sparse_matrix_hdf5(filename, sparse_matrix, dataset_name = 'sparse_matrix', format_ = 'csr'):
+	"""
+	Writes a sparse matrix to an HDF5 file in a specified format.
+
+	Args:
+		filename (str): Name of the output file
+		sparse_matrix (scipy.sparse matrix): Sparse matrix to be written to the file.
+		dataset_name (str, optional): Name of the HDF5 dataset. Defaults to 'sparse_matrix'.
+		format_ (str, optional): The format of the sparse matrix. Currently only supports 'csr' (Compressed Sparse Row). 
+	"""
+	if format_ == 'csr':
+		with h5py.File(filename, 'w') as f:
+			g = f.create_group(dataset_name)
+			g.create_dataset('data', data = sparse_matrix.data)
+			g.create_dataset('indices', data = sparse_matrix.indices)
+			g.create_dataset('indptr', data = sparse_matrix.indptr)
+			g.create_dataset('shape', data = np.array(sparse_matrix.shape))
+			g.attrs['format'] = 'csr'
+	else:
+		raise ValueError(f'Unsupported sparse matrix format: {format_}')
+
+
+def read_sparse_matrix_hdf5(filename, dataset_name = 'sparse_matrix'):
+	"""
+	Reads a sparse matrix from an HDF5 file.
+
+	Args:
+		filename (str): Name of the output file
+		dataset_name (str, optional): Name of the dataset containing the matrix. Defaults to 'sparse_matrix'.
+
+	Returns:
+		scipy.sparse matrix: The sparse matrix read from the file.
+	"""
+	with h5py.File(filename, 'r') as f:
+		g = f[dataset_name]
+		data = g['data'][:]
+		indices = g['indices'][:]
+		indptr = g['indptr'][:]
+		shape = tuple(g['shape'][:])
+		format_ = g.attrs['format']
+		if format_ == 'csr':
+			return csr_matrix((data, indices, indptr), shape = shape)
+		else:
+			raise ValueError(f'Unsupported sparse matrix format: {format_}')
+
+
+def data_dump(bf, rho, rho_pred, params_dict):
+	"""
+	Writes data and config to file for later use.
+	"""
+	np.savez('sample_data.npz', rho = rho, rho_pred = rho_pred, bf = bf)
+	json.dump('sample_config.json', params_dict, encoding = 'utf-8')
+
+
+def save_as_frozen_graph(model, saved_model_dir):
+	"""
+	Converts a TensorFlow model into a 'frozen' SavedModel format.
+
+	Args:
+		model: TensorFlow model to be frozen.
+		saved_model_dir (str): The directory path where the frozen model will be saved.
+
+	Returns:
+		TFModelServer object: An instance of the TFModelServer class, which can be used for serving the model.
+
+	The function converts model variables to constants and is required for converting the model to the intermediate
+	 representation (IR) used by openvino.
+	"""
+	# Create input specifications for the model
+	input_specs = [tf.TensorSpec([1] + model.inputs[i].shape[1:].as_list(), model.inputs[i].dtype) for i in
+				   range(len(model.inputs))]
+
+	# Create a concrete function from the model
+	full_model = tf.function(lambda x: model(x))
+	full_model = full_model.get_concrete_function(input_specs)
+
+	# Convert the model to a frozen function
+	frozen_func = convert_variables_to_constants_v2(full_model)
+
+	# Define a new module with a method that has a `@tf.function` decorator with input signatures
+	class TFModelServer(tf.Module):
+		def __init__(self, frozen_func):
+			super().__init__()
+			self.frozen_func = frozen_func
+
+		@tf.function(input_signature = input_specs)
+		def serve(self, *args):
+			return self.frozen_func(*args)
+
+	# Create an instance of TFModelServer with the frozen function
+	model_server = TFModelServer(frozen_func)
+
+	# Save the module as a SavedModel
+	tf.saved_model.save(model_server, saved_model_dir, signatures = {"serving_default": model_server.serve})
+
+	return model_server
+
+
+def log_and_print(message):
+	print(message)
+	with open('log.txt', 'a') as log_file:
+		log_file.write(message + '\n')

stnn/utils/network_visualization.py

@@ -0,0 +1,121 @@
+import pydot
+import re
+from keras.models import Model
+from keras.layers import Layer, InputLayer
+from pygments.lexers import graphviz
+
+
+# May be necessary to manually add Graphviz to PATH, e.g.
+# import os
+# os.environ["PATH"] += os.pathsep + r'C:\Program Files\Graphviz\bin'
+
+def visualize_model(model, layer_labels = None, layer_colors = None, groupings = None, exclude_input_layer = False,
+					verbose = False, output_filename = 'model_graph.png'):
+	"""
+	Creates a visual graph of a keras model. There is an option to group certain layers into subgraphs
+	(argument 'groupings').
+
+	Args:
+		model: A Keras Model instance
+		layer_labels (optional): List of labels for each layer. Defaults to layer names.
+		layer_colors (optional): List of colors for each layer. Defaults to white for all layers.
+		groupings (optional): Dictionary specifying groups of layers. Each key is a group name,
+							  and its value is a list of layer names belonging to that group.
+		exclude_input_layer (optional): Boolean indicating whether to exclude the input layer from the graph.
+		verbose (boolean, optional): Whether to print verbose output. Defaults to False.
+		output_filename (optional): name of the output file for saving the generated graph.
+
+	Output:
+		Image file with name 'output_filename'.
+	"""
+	if not isinstance(model, Model):
+		raise ValueError("model should be a Keras model instance")
+	num_layers = len(model.layers)
+
+	# Default labels and colors if not provided
+	if not layer_labels:
+		layer_labels = [layer.name for layer in model.layers]
+	if not layer_colors:
+		default_color = 'white'
+		layer_colors = [default_color] * num_layers
+
+	# Create a directed graph
+	graph = pydot.Dot(graph_type = 'digraph', rankdir = 'LR')
+
+	# Create nodes for each layer and add to subgraphs if specified
+	subgraphs = {}
+	layer_id_map = {}
+	for i, layer in enumerate(model.layers):
+		# Exclude the input layer if specified
+		if exclude_input_layer and isinstance(layer, InputLayer):
+			continue
+
+		# Create a node for the layer
+		layer_id = str(id(layer))
+		layer_id_map[layer] = layer_id
+		label = layer_labels[i]
+		color = layer_colors[i]
+
+		node = pydot.Node(layer_id, label = label, style = 'filled', fillcolor = color, shape = 'box')
+
+		# Check for groupings and add the node to the appropriate subgraph or main graph
+		group_name = None
+		if groupings:
+			for group, members in groupings.items():
+				if layer.name in members:
+					group_name = group
+					break
+
+		if group_name:
+			if group_name not in subgraphs:
+				subgraph = pydot.Cluster(group_name, label = group_name, style = 'dashed', fontsize = 24)
+				subgraphs[group_name] = subgraph
+			subgraphs[group_name].add_node(node)
+		else:
+			graph.add_node(node)
+
+	# Add subgraphs to the main graph
+	for subgraph in subgraphs.values():
+		graph.add_subgraph(subgraph)
+
+	# Add edges based on layer connections
+	for layer in model.layers:
+		if exclude_input_layer and isinstance(layer, InputLayer):
+			continue
+		# Handle custom or non-standard layers
+		if hasattr(layer, '_inbound_nodes'):
+			inbound_nodes = layer._inbound_nodes
+		else:
+			# If the layer doesn't have '_inbound_nodes', skip edge creation
+			continue
+
+		inbound_layers = []
+		for inbound_node in inbound_nodes:
+			inbound_layers = inbound_node.inbound_layers
+			if not isinstance(inbound_layers, list):
+				inbound_layers = [inbound_layers]
+
+		for inbound_node in inbound_nodes:
+			for inbound_layer in inbound_layers:
+				if isinstance(inbound_layer, Layer) and inbound_layer in layer_id_map:
+					src_id = layer_id_map[inbound_layer]
+					dest_id = layer_id_map[layer]
+					if (re.search('sequential', inbound_layer.name, flags = re.IGNORECASE) or
+							re.search(r'operators__.getitem_[0-9]+$', inbound_layer.name, flags = re.IGNORECASE)):
+						graph.add_edge(pydot.Edge(src_id, dest_id, style = 'invis'))
+					else:
+						graph.add_edge(pydot.Edge(src_id, dest_id))
+					if verbose:
+						print(f"Added edge from {inbound_layer.name} to {layer.name}")
+
+	graph.set_graph_defaults(sep = '+125,125')
+	try:
+		graph.write_png(output_filename)
+	except FileNotFoundError as e:
+		print(f'\nFailed to create network visualization using pydot and graphviz. Pleasure ensure that '
+			  'the output filename is valid, and graphviz is installed and included in the system PATH variable. '
+			  f'Original error: {e}')
+	except Exception as e:
+		print(f'\nFailed to create network visualization using pydot and graphviz. Original error: {e}')
+	else:
+		print(f'Model visualization saved to {output_filename}')

stnn/utils/plotting.py

@@ -0,0 +1,93 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from matplotlib.cm import get_cmap
+
+
+def plot_comparison(system, bf, rho, rho_pred, fontscale = 1, output_filename = 'comparison.png',
+					wspace = -0.1, hspace = 0.5):
+	"""
+	Side-by-side comparison of contour plots for arguments 'rho' and 'rho_pred'.
+	Also plots the boundary data (argument 'bf') as a contour plot.
+
+	Args:
+		system: PDESystem object
+		bf (np.ndarray): 2D array representing the boundary data on the (x2, x3) grid.
+		rho (np.ndarray): 2D array, known rho
+		rho_pred (np.ndarray): 2D array, predicted rho
+		fontscale (int, optional): A scaling factor for the font size in plots.
+		output_filename (str, optional): name of the output file
+		wspace (float, optional): horizontal spacing between subplots
+		hspace (float, optional): vertical spacing between subplots
+
+	Returns:
+		Does not return a value. The figure is saved to a file with name given by 'output_filename'. The default
+		is 'comparison.png'.
+	"""
+	# System parameters
+	ell = system.params['ell']
+	a2 = system.a2
+	e = system.params['eccentricity']
+	nx1, nx2, nx3 = system.params['nx1'], system.params['nx2'], system.params['nx3']
+
+	# Get x, y grids from 'PDESystem' object
+	x, y = system.get_xy_grids()
+
+	# Relative error
+	err = np.linalg.norm(rho - rho_pred)
+	rel_err = err / np.linalg.norm(rho)
+
+	# wrap around values for continuity
+	rho = np.append(rho, rho[:, 0:1], axis = 1)
+	rho_pred = np.append(rho_pred, rho_pred[:, 0:1], axis = 1)
+
+	plots = [rho, rho_pred]
+
+	# Color bar limits
+	vmin = np.nanmin(rho_pred[:, :])
+	vmax = np.nanmax(rho_pred[:, :])
+	vmin = min(vmin, np.nanmin(rho[:, :]))
+	vmax = max(vmax, np.nanmax(rho[:, :]))
+
+	# Figure layout
+	cbar_coords = (0.89, 0.11, 0.03, 0.77)
+	fig = plt.figure(figsize = (24, 24))
+	gs = fig.add_gridspec(11, 2, hspace = hspace, wspace = wspace)
+	axs = [fig.add_subplot(gs[:11, 0]), fig.add_subplot(gs[:5, 1]), fig.add_subplot(gs[6:11, 1])]
+
+	# boundary data contour plot
+	bf_plot = np.nan * np.ones((2 * nx2, nx3))
+	bf_plot[:nx2, :nx3 // 2] = bf[:nx2, :]
+	bf_plot[nx2:, nx3 // 2:] = bf[nx2:, :]
+	im = axs[0].imshow(bf_plot)
+	axs[0].set_yticks([1, nx2 // 4, nx2 // 2 - 1, nx2 // 2 + 1, 3 * nx2 // 4, nx2])
+	axs[0].set_yticklabels(['-pi', '0', 'pi', '', '0', 'pi'])
+	axs[0].set_xticks([1, nx3 // 2, nx3])
+	axs[0].set_xticklabels(['0', 'pi', '2pi'])
+	axs[0].set_title('Boundary data\n\n', fontsize = fontscale * 48, y = 0.91)
+	for label in axs[0].get_xticklabels() + axs[0].get_yticklabels():
+		label.set_fontsize(fontscale * 48)
+	divider = make_axes_locatable(axs[0])
+	cax = divider.append_axes('bottom', size = "3%", pad = 1.5)
+	cb = fig.colorbar(im, cax = cax, orientation = 'horizontal')
+	cb.ax.tick_params(labelsize = fontscale * 48)
+	cax.xaxis.set_ticks_position('bottom')
+
+	# rho(x, y) contour plots
+	titles = ['Direct Solution', 'Tensor network']
+	for i, ax in enumerate(axs[1:]):
+		z = plots[i]
+		im = ax.contourf(x, y, z, levels = np.linspace(vmin, vmax, 100), cmap = get_cmap('hsv'))
+		ax.set_title(titles[i], fontsize = fontscale * 48)
+		for label in ax.get_xticklabels() + ax.get_yticklabels():
+			label.set_fontsize(fontscale * 32)
+		ax.set_aspect(1.0)
+
+	cbar_ax = fig.add_axes(cbar_coords)
+	cb = fig.colorbar(im, cax = cbar_ax, pad = 0.05)
+	cb.ax.tick_params(labelsize = fontscale * 32)
+
+	suptitle = f'ell = {ell:.3f}; a2 = {a2:.3f}; e = {e:.3f}; Relative error: {rel_err:.3f}'
+	plt.suptitle(suptitle, fontsize = fontscale * 48, x = 0.1, y = 0.97, horizontalalignment = 'left')
+	plt.savefig(output_filename)
+	plt.close()

stnn/utils/stats.py

@@ -0,0 +1,27 @@
+import numpy as np
+
+
+def get_stats(rho, rho_pred, output_filename = 'stats.npz'):
+	"""
+	Calculates statistical metrics for model predictions and saves them to a file.
+
+	This function computes the normalized loss for each instance in the dataset, identifies the maximum loss and its 
+	index, and calculates the average loss. These statistics are then saved to an NPZ file.
+
+	Args:
+		rho (numpy.ndarray): True values.
+		rho_pred (numpy.ndarray): Predicted values.
+		output_filename (str, optional): The name of the file where the statistics will be saved. Defaults to 'stats.npz'.
+	"""
+	if rho.shape != rho_pred.shape:
+		raise ValueError('rho and rho_pred must have the same shape.')
+
+	y_true_flattened = rho.reshape(rho.shape[0], -1)
+	y_pred_flattened = rho_pred.reshape(rho_pred.shape[0], -1)
+	loss = np.linalg.norm(y_true_flattened - y_pred_flattened, axis = 1) / np.linalg.norm(y_true_flattened, axis = 1)
+	max_loss = np.max(loss)
+	max_loss_index = np.argmax(loss)
+	print(f'Maximum loss: {max_loss}')
+	print(f'Index of instance with max loss: {max_loss_index}')
+	print(f'Average loss: {np.average(loss)}')
+	np.savez(output_filename, max_loss = max_loss, avg_loss = np.average(loss), N = loss.shape[0])