Add space-time simulator Python script

2ed2c55 almost 2 years ago

6.15 kB

import numpy as np
from typing import Tuple, Union

# Simulate your friend's typing_ and backend modules
# In your real code, replace these with your actual imports.
Number = Union[float, int]
Tensorlike = np.ndarray
bd = np

# Constants (if needed)
# from . import constants as const
# For now we won't define constants as none were used explicitly.

def curl_E(E: Tensorlike) -> Tensorlike:
    """
    Transforms an E-type field into an H-type field by performing a curl operation.
    E is located on edges (integer gridpoints), result is on faces (half-integer grid points).
    """
    curl = bd.zeros(E.shape)

    # Following the original indexing logic
    curl[:, :-1, :, 0] += E[:, 1:, :, 2] - E[:, :-1, :, 2]
    curl[:, :, :-1, 0] -= E[:, :, 1:, 1] - E[:, :, :-1, 1]

    curl[:, :, :-1, 1] += E[:, :, 1:, 0] - E[:, :, :-1, 0]
    curl[:-1, :, :, 1] -= E[1:, :, :, 2] - E[:-1, :, :, 2]

    curl[:-1, :, :, 2] += E[1:, :, :, 1] - E[:-1, :, :, 1]
    curl[:, :-1, :, 2] -= E[:, 1:, :, 0] - E[:, :-1, :, 0]

    return curl

def curl_edge_to_face(v: Tensorlike) -> Tensorlike:
    """
    Compute the curl of a vector field v on edges to get a field on faces.
    """
    return curl_E(v)

def curl_H(H: Tensorlike) -> Tensorlike:
    """
    Transforms an H-type field (faces) into an E-type field (edges) by performing a curl operation.
    """
    curl = bd.zeros(H.shape)

    curl[:, 1:, :, 0] += H[:, 1:, :, 2] - H[:, :-1, :, 2]
    curl[:, :, 1:, 0] -= H[:, :, 1:, 1] - H[:, :, :-1, 1]

    curl[:, :, 1:, 1] += H[:, :, 1:, 0] - H[:, :, :-1, 0]
    curl[1:, :, :, 1] -= H[1:, :, :, 2] - H[:-1, :, :, 2]

    curl[1:, :, :, 2] += H[1:, :, :, 1] - H[:-1, :, :, 1]
    curl[:, 1:, :, 2] -= H[:, 1:, :, 0] - H[:, :-1, :, 0]

    return curl

def curl_face_to_edge(A: Tensorlike) -> Tensorlike:
    """
    Compute the curl of a vector field A on faces, resulting in a field on edges.
    """
    return curl_H(A)

def div(v: Tensorlike) -> Tensorlike:
    """
    Compute the divergence of a vector field v defined on edges.
    The result is a scalar field on faces.
    """
    # Based on your your old code
    div_v = bd.zeros((v.shape[0], v.shape[1], v.shape[2], 1))

    # x-component
    div_v[1:-1, :, :, 0] += (v[1:-1, :, :, 0] - v[:-2, :, :, 0])
    div_v[:-2, :, :, 0] -= (v[1:-1, :, :, 0] - v[2:, :, :, 0])

    # y-component
    div_v[:, 1:-1, :, 0] += (v[:, 1:-1, :, 1] - v[:, :-2, :, 1])
    div_v[:, :-2, :, 0]  -= (v[:, 1:-1, :, 1] - v[:, 2:, :, 1])

    # z-component
    div_v[:, :, 1:-1, 0] += (v[:, :, 1:-1, 2] - v[:, :, :-2, 2])
    div_v[:, :, :-2, 0]  -= (v[:, :, 1:-1, 2] - v[:, :, 2:, 2])

    return div_v

def grad(p: Tensorlike) -> Tensorlike:
    """
    Compute the gradient of a scalar field p located on grid points.
    p has shape (Nx+1, Ny+1, Nz+1, 1).
    Returns (Nx+1, Ny+1, Nz+1, 3).
    """
    grad = bd.zeros((p.shape[0], p.shape[1], p.shape[2], 3))
    
    grad[:-1, :, :, 0] = (p[1:, :, :, 0] - p[:-1, :, :, 0])
    grad[:, :-1, :, 1] = (p[:, 1:, :, 0] - p[:, :-1, :, 0])
    grad[:, :, :-1, 2] = (p[:, :, 1:, 0] - p[:, :, :-1, 0])
    
    return grad

############################################
# Additional functionalities and example
############################################

def laplacian_vector_field(v: Tensorlike) -> Tensorlike:
    """
    Compute the vector Laplacian Δv using the provided curl and div.
    Δv = ∇(∇·v) - ∇×(∇×v)
    To do this correctly, we might need v and p fields on compatible staggered grids.
    Here we assume the same indexing as your friend's code.
    """
    # As an example, let's just apply curl and div in a way similar to the previous code
    # Note: The original code doesn't provide dx, dy, dz, we assume uniform spacing of 1.
    div_v = div(v)              # scalar field
    grad_div_v = grad(div_v)    # vector field
    c = curl_edge_to_face(v)    # curl moves edges->faces
    curl_c = curl_face_to_edge(c) # curl moves faces->edges
    return grad_div_v - curl_c  # Δv = ∇(∇·v) - ∇×(∇×v)

def apply_second_order_operator(v: Tensorlike, k: float = 1.0) -> Tensorlike:
    """
    Apply d/dt = -k Δ to a vector field v on edges.
    Here Δv = laplacian_vector_field(v).
    """
    lap_v = laplacian_vector_field(v)
    return -k * lap_v

# Boundary conditions
def apply_dirichlet_condition(field: np.ndarray, value: float = 0.0):
    field[0, :, :, :] = value
    field[-1, :, :, :] = value
    field[:, 0, :, :] = value
    field[:, -1, :, :] = value
    field[:, :, 0, :] = value
    field[:, :, -1, :] = value

def apply_neumann_condition(field: np.ndarray):
    field[0, :, :, :] = field[1, :, :, :]
    field[-1, :, :, :] = field[-2, :, :, :]
    field[:, 0, :, :] = field[:, 1, :, :]
    field[:, -1, :, :] = field[:, -2, :, :]
    field[:, :, 0, :] = field[:, :, 1, :]
    field[:, :, -1, :] = field[:, :, -2, :]

def apply_periodic_condition(field: np.ndarray):
    field[0, :, :, :] = field[-2, :, :, :] 
    field[-1, :, :, :] = field[1, :, :, :]
    field[:, 0, :, :] = field[:, -2, :, :]
    field[:, -1, :, :] = field[:, 1, :, :]
    field[:, :, 0, :] = field[:, :, -2, :]
    field[:, :, -1, :] = field[:, :, 1, :]

if __name__ == "__main__":
    # Example usage:
    # Suppose Nx, Ny, Nz are the number of cells in each direction
    # The given code suggests fields on staggered grids. Adjust sizes as needed.
    Nx, Ny, Nz = 20, 20, 20
    v_field = bd.random.randn(Nx, Ny, Nz, 3)  # A random vector field
    
    # Choose a boundary type
    boundary_type = 'dirichlet'  # 'dirichlet', 'neumann', 'periodic'
    dt = 0.1
    k = 1.0
    num_timesteps = 5

    for _ in range(num_timesteps):
        dv_dt = apply_second_order_operator(v_field, k)
        v_field = v_field + dt * dv_dt

        # Apply chosen boundary condition
        if boundary_type == 'dirichlet':
            apply_dirichlet_condition(v_field, value=0.0)
        elif boundary_type == 'neumann':
            apply_neumann_condition(v_field)
        elif boundary_type == 'periodic':
            apply_periodic_condition(v_field)

    print("Final field shape:", v_field.shape)
    print("Simulation complete.")

	import numpy as np
	from typing import Tuple, Union

	# Simulate your friend's typing_ and backend modules
	# In your real code, replace these with your actual imports.
	Number = Union[float, int]
	Tensorlike = np.ndarray
	bd = np

	# Constants (if needed)
	# from . import constants as const
	# For now we won't define constants as none were used explicitly.

	def curl_E(E: Tensorlike) -> Tensorlike:
	"""
	Transforms an E-type field into an H-type field by performing a curl operation.
	E is located on edges (integer gridpoints), result is on faces (half-integer grid points).
	"""
	curl = bd.zeros(E.shape)

	# Following the original indexing logic
	curl[:, :-1, :, 0] += E[:, 1:, :, 2] - E[:, :-1, :, 2]
	curl[:, :, :-1, 0] -= E[:, :, 1:, 1] - E[:, :, :-1, 1]

	curl[:, :, :-1, 1] += E[:, :, 1:, 0] - E[:, :, :-1, 0]
	curl[:-1, :, :, 1] -= E[1:, :, :, 2] - E[:-1, :, :, 2]

	curl[:-1, :, :, 2] += E[1:, :, :, 1] - E[:-1, :, :, 1]
	curl[:, :-1, :, 2] -= E[:, 1:, :, 0] - E[:, :-1, :, 0]

	return curl

	def curl_edge_to_face(v: Tensorlike) -> Tensorlike:
	"""
	Compute the curl of a vector field v on edges to get a field on faces.
	"""
	return curl_E(v)

	def curl_H(H: Tensorlike) -> Tensorlike:
	"""
	Transforms an H-type field (faces) into an E-type field (edges) by performing a curl operation.
	"""
	curl = bd.zeros(H.shape)

	curl[:, 1:, :, 0] += H[:, 1:, :, 2] - H[:, :-1, :, 2]
	curl[:, :, 1:, 0] -= H[:, :, 1:, 1] - H[:, :, :-1, 1]

	curl[:, :, 1:, 1] += H[:, :, 1:, 0] - H[:, :, :-1, 0]
	curl[1:, :, :, 1] -= H[1:, :, :, 2] - H[:-1, :, :, 2]

	curl[1:, :, :, 2] += H[1:, :, :, 1] - H[:-1, :, :, 1]
	curl[:, 1:, :, 2] -= H[:, 1:, :, 0] - H[:, :-1, :, 0]

	return curl

	def curl_face_to_edge(A: Tensorlike) -> Tensorlike:
	"""
	Compute the curl of a vector field A on faces, resulting in a field on edges.
	"""
	return curl_H(A)

	def div(v: Tensorlike) -> Tensorlike:
	"""
	Compute the divergence of a vector field v defined on edges.
	The result is a scalar field on faces.
	"""
	# Based on your your old code
	div_v = bd.zeros((v.shape[0], v.shape[1], v.shape[2], 1))

	# x-component
	div_v[1:-1, :, :, 0] += (v[1:-1, :, :, 0] - v[:-2, :, :, 0])
	div_v[:-2, :, :, 0] -= (v[1:-1, :, :, 0] - v[2:, :, :, 0])

	# y-component
	div_v[:, 1:-1, :, 0] += (v[:, 1:-1, :, 1] - v[:, :-2, :, 1])
	div_v[:, :-2, :, 0] -= (v[:, 1:-1, :, 1] - v[:, 2:, :, 1])

	# z-component
	div_v[:, :, 1:-1, 0] += (v[:, :, 1:-1, 2] - v[:, :, :-2, 2])
	div_v[:, :, :-2, 0] -= (v[:, :, 1:-1, 2] - v[:, :, 2:, 2])

	return div_v

	def grad(p: Tensorlike) -> Tensorlike:
	"""
	Compute the gradient of a scalar field p located on grid points.
	p has shape (Nx+1, Ny+1, Nz+1, 1).
	Returns (Nx+1, Ny+1, Nz+1, 3).
	"""
	grad = bd.zeros((p.shape[0], p.shape[1], p.shape[2], 3))

	grad[:-1, :, :, 0] = (p[1:, :, :, 0] - p[:-1, :, :, 0])
	grad[:, :-1, :, 1] = (p[:, 1:, :, 0] - p[:, :-1, :, 0])
	grad[:, :, :-1, 2] = (p[:, :, 1:, 0] - p[:, :, :-1, 0])

	return grad

	############################################
	# Additional functionalities and example
	############################################

	def laplacian_vector_field(v: Tensorlike) -> Tensorlike:
	"""
	Compute the vector Laplacian Δv using the provided curl and div.
	Δv = ∇(∇·v) - ∇×(∇×v)
	To do this correctly, we might need v and p fields on compatible staggered grids.
	Here we assume the same indexing as your friend's code.
	"""
	# As an example, let's just apply curl and div in a way similar to the previous code
	# Note: The original code doesn't provide dx, dy, dz, we assume uniform spacing of 1.
	div_v = div(v) # scalar field
	grad_div_v = grad(div_v) # vector field
	c = curl_edge_to_face(v) # curl moves edges->faces
	curl_c = curl_face_to_edge(c) # curl moves faces->edges
	return grad_div_v - curl_c # Δv = ∇(∇·v) - ∇×(∇×v)

	def apply_second_order_operator(v: Tensorlike, k: float = 1.0) -> Tensorlike:
	"""
	Apply d/dt = -k Δ to a vector field v on edges.
	Here Δv = laplacian_vector_field(v).
	"""
	lap_v = laplacian_vector_field(v)
	return -k * lap_v

	# Boundary conditions
	def apply_dirichlet_condition(field: np.ndarray, value: float = 0.0):
	field[0, :, :, :] = value
	field[-1, :, :, :] = value
	field[:, 0, :, :] = value
	field[:, -1, :, :] = value
	field[:, :, 0, :] = value
	field[:, :, -1, :] = value

	def apply_neumann_condition(field: np.ndarray):
	field[0, :, :, :] = field[1, :, :, :]
	field[-1, :, :, :] = field[-2, :, :, :]
	field[:, 0, :, :] = field[:, 1, :, :]
	field[:, -1, :, :] = field[:, -2, :, :]
	field[:, :, 0, :] = field[:, :, 1, :]
	field[:, :, -1, :] = field[:, :, -2, :]

	def apply_periodic_condition(field: np.ndarray):
	field[0, :, :, :] = field[-2, :, :, :]
	field[-1, :, :, :] = field[1, :, :, :]
	field[:, 0, :, :] = field[:, -2, :, :]
	field[:, -1, :, :] = field[:, 1, :, :]
	field[:, :, 0, :] = field[:, :, -2, :]
	field[:, :, -1, :] = field[:, :, 1, :]

	if __name__ == "__main__":
	# Example usage:
	# Suppose Nx, Ny, Nz are the number of cells in each direction
	# The given code suggests fields on staggered grids. Adjust sizes as needed.
	Nx, Ny, Nz = 20, 20, 20
	v_field = bd.random.randn(Nx, Ny, Nz, 3) # A random vector field

	# Choose a boundary type
	boundary_type = 'dirichlet' # 'dirichlet', 'neumann', 'periodic'
	dt = 0.1
	k = 1.0
	num_timesteps = 5

	for _ in range(num_timesteps):
	dv_dt = apply_second_order_operator(v_field, k)
	v_field = v_field + dt * dv_dt

	# Apply chosen boundary condition
	if boundary_type == 'dirichlet':
	apply_dirichlet_condition(v_field, value=0.0)
	elif boundary_type == 'neumann':
	apply_neumann_condition(v_field)
	elif boundary_type == 'periodic':
	apply_periodic_condition(v_field)

	print("Final field shape:", v_field.shape)
	print("Simulation complete.")