DIII-D / GITRmProject

Apache License 2.0
0 stars 0 forks source link

generate a dimes mesh with python #5

Closed jguterl closed 1 month ago

jguterl commented 1 month ago
import matplotlib.pyplot as plt
import numpy as np
import gmsh
import math

class Point3d(object):
    def __init__(self, x=0, y=0, z=0):
        if type(x) == list:
            # if receives a list
            try:
                assert len(x) == 3
                self.coords = np.array(x+[1], dtype=np.float64)
            except:
                raise ValueError(
                    'To create a Point3d with a list it must have length 3')
        elif type(x) == tuple:
            # if receives a tuple
            try:
                assert len(x) == 3
                self.coords = np.array(list(x)+[1], dtype=np.float64)
            except:
                raise ValueError(
                    'To create a Point3d with a tuple it must have length 3')
        elif type(x) == np.ndarray:
            try:
                assert x.ndim == 1 and x.size == 4 and x[3] == 1
                self.coords = np.array(x, dtype=np.float64)
            except:
                print("ndim", x.ndim)
                print("size", x.size)
                print("[3]", x[3])
                raise ValueError(
                    'To create a Point3d with an np.array it must have ndim 1, size 4, and the last element must be 1')
        elif type(x) == Point3d:
            self.coords = np.array(x.coords, dtype=np.float64)
        else:
            self.coords = np.array([x, y, z, 1], dtype=np.float64)

        self.iter = 0

    def copy(self):
        return Point3d(self.coords)

    def __getattr__(self, name):
        """ When the user asks for attributes x, y, or z, we return
            coords[0], [1], and [2] """
        if name == 'x':
            return self.coords[0]
        elif name == 'y':
            return self.coords[1]
        elif name == 'z':
            return self.coords[2]

    def __getitem__(self, key):
        return self.coords[key]

    def __setattr__(self, name, value):
        """ For x, y, and z sets coords[0], [1], and [2].
            For everything else does the normal set. """
        if name == 'x':
            self.coords[0] = value
        elif name == 'y':
            self.coords[1] = value
        elif name == 'z':
            self.coords[2] = value
        else:
            self.__dict__[name] = value

    def __setitem__(self, key, value):
        self.coords[key] = value

    def __repr__(self):
        return "Point3d({0}, {1}, {2})".format(self.x, self.y, self.z)

    def __str__(self):
        return "({0}, {1}, {2})".format(self.x, self.y, self.z)

    def __add__(self, other):
        x = self.x + other.x
        y = self.y + other.y
        z = self.z + other.z
        return Point3d(x, y, z)

    def __sub__(self, other):
        x = self.x - other.x
        y = self.y - other.y
        z = self.z - other.z
        return Point3d(x, y, z)

    def __mul__(self, other):
        x = self.x * other
        y = self.y * other
        z = self.z * other
        return Point3d(x, y, z)

    def dot(self, other):
        return self.coords.dot(other.coords)

    def cross(self, other):
        a = np.array([self.x, self.y, self.z])
        b = np.array([other.x, other.y, other.z])
        c = np.cross(a, b)
        return Point3d(list(c))

    def __pos__(self):
        return self

    def __neg__(self):
        return self * -1

    def __iter__(self):
        return self

    def next(self):
        if self.iter >= self.coords.size-1:
            self.iter = 0
            raise StopIteration
        else:
            self.iter += 1
            return self.coords[self.iter-1]

    def length(self):
        return math.sqrt(self.x**2 + self.y**2 + self.z**2)

    def translate(self, vector):
        self.x += vector.x
        self.y += vector.y
        self.z += vector.z

    def rotateX(self, radians):
        c = np.cos(radians)
        s = np.sin(radians)
        xmat = np.array([[1, 0, 0, 0],
                        [0, c, -s, 0],
                        [0, s, c, 0],
                        [0, 0, 0, 1]])
        self.coords = xmat.dot(self.coords)

    def rotateY(self, radians):
        c = np.cos(radians)
        s = np.sin(radians)
        ymat = np.array([[c, 0, s, 0],
                        [0, 1, 0, 0],
                        [-s, 0, c, 0],
                        [0, 0, 0, 1]])
        self.coords = ymat.dot(self.coords)

    def rotateZ(self, radians):
        c = np.cos(radians)
        s = np.sin(radians)
        zmat = np.array([[c, -s, 0, 0],
                        [s, c, 0, 0],
                        [0, 0, 1, 0],
                        [0, 0, 0, 1]])
        self.coords = zmat.dot(self.coords)

# from .surface_elements import *
occ = gmsh.model.occ
Point = Point3d

def dump_mesh(fn):
    gmsh.write(f"{fn}.msh")
    write_partition_file(fn)
    # write_surface_elements(fn)

def write_partition_file(fn):
    N_elem = len(gmsh.model.mesh.get_elements(3)[1][0])
    with open(f"{fn}.part", "w") as f:
        for i in range(N_elem):
            f.write("0\n")

# def write_surface_elements(fn):
#     elements = SurfaceElements()
#     elements.import_mesh_elements_from_model(gmsh.model)
#     elements.save(f"{fn}.pkl")

# for r_large in [0.002, 0.003, 0.004, 0.005, 0.0075, 0.01, 0.0125]:
for r_large in [0.002]:
    # Initialize gmsh session

    gmsh.initialize()

    # Coordinate system: (R,y,Z)
    # Reference plan Z
    Z_ref = -1.25

    # Define DiMES cap surface
    c_DiMES = Point(1.485, 0.0, Z_ref)
    r_DiMES = 0.025

    # Define small dot on DiMES cap surface with respect to DiMES center
    c_smallDot = Point(-0.015, 0.0, 0.0) + c_DiMES
    c_largeDot = Point(0.00, 0.0, 0.0) + c_DiMES
    r_smallDot = 0.0005
    r_largeDot = r_large

    # Simulation box
    l_box = (r_DiMES + 0.015) * 2
    z_box = l_box
    v_box = c_DiMES - Point(l_box/2, l_box/2, 0.0)
    # Tag

    # Create 2D surfaces
    box = occ.addBox(v_box.x, v_box.y, v_box.z, l_box, l_box, z_box)
    box_faces = occ.get_entities(2)
    box_verteces = occ.get_entities(1)
    DiMES0 = gmsh.model.occ.addDisk(
        c_DiMES.x, c_DiMES.y, c_DiMES.z, r_DiMES, r_DiMES)

    occ.remove([(3, box)])

    def make_dot(c, r):
        return occ.addDisk(c.x, c.y, c.z, r, r)

    largeDot = make_dot(c_largeDot, r_largeDot)
    # smallDot = make_dot(c_smallDot, r_smallDot)
    # print(largeDot, smallDot)
    # DiMES = gmsh.model.occ.cut(
    #   [(2, DiMES0)], [(2, smallDot), (2, largeDot)], removeTool=False, removeObject=False)[0][0][1]
    DiMES = gmsh.model.occ.cut(
        [(2, DiMES0)], [(2, largeDot)], removeTool=False, removeObject=False)[0][0][1]
    print(DiMES)
    gmsh.model.occ.remove(gmsh.model.get_entities(1))
    gmsh.model.occ.remove(gmsh.model.get_entities(0))
    id_box_bottom = 5
    bottom_box = gmsh.model.occ.cut([(2, id_box_bottom)], [(
        2, DiMES0)], removeTool=True, removeObject=True)[0][0][1]

    # surfs = [DiMES, largeDot, smallDot, bottom_box] + [v_[1] for v_ in box_faces if (v_[0] == 2 and v_[1] != bottom_box)]
    surfs = [DiMES, largeDot,  bottom_box] + [v_[1]
                                              for v_ in box_faces if (v_[0] == 2 and v_[1] != bottom_box)]
    print(surfs)
    surface_loop = gmsh.model.occ.add_surface_loop(surfs)
    vol = gmsh.model.occ.add_volume([surface_loop])
    gmsh.model.occ.remove([(2, surface_loop)])
    gmsh.model.occ.remove(gmsh.model.get_entities(2))
    gmsh.model.occ.remove(gmsh.model.get_entities(1))
    gmsh.model.occ.remove(gmsh.model.get_entities(0))

    # # Synchronize necessary before mesh setup and generation
    gmsh.model.occ.synchronize()

    # # Set number of elements on the boundary of each dots and DiMES cap
    gmsh.option.setNumber("Mesh.MeshSizeFromCurvature", 1)
    gmsh.option.setNumber("Mesh.MinimumElementsPerTwoPi", 20)

    # # Prevent very small elements in small dots
    gmsh.option.setNumber("Mesh.MeshSizeMin", 0.0003)
    gmsh.option.setNumber("Mesh.MeshSizeMax", 0.001)
    # Launch the GUI to see the results:

    # for i in range(15):
    #     gmsh.model.occ.remove([(0, i+1)])
    # Generate 2D mesh
    # gmsh.fltk.run()
    mesh = gmsh.model.mesh.generate(3)
    gmsh.fltk.run()

    # Write mesh into a meshio format
    fn = f'small_large_dots_DiMES_uniform_r_{r_large}'
    gmsh.model.mesh.remove_duplicate_elements()
    gmsh.model.mesh.remove_duplicate_nodes()
    dump_mesh(fn)

    # Close gmsh session
    gmsh.finalize()
LC-GA commented 1 month ago
#%%
""" miscellaneous functions """

def rectangle_def(x, y, z, width, height):
    # Create points for the rectangle corners
    p1 = gmsh.model.occ.addPoint(x, y, z)          # Bottom-left corner
    p2 = gmsh.model.occ.addPoint(x + width, y, z)  # Bottom-right corner
    p3 = gmsh.model.occ.addPoint(x + width, y + height, z)  # Top-right corner
    p4 = gmsh.model.occ.addPoint(x, y + height, z)  # Top-left corner

    # Create lines for the rectangle edges
    l1 = gmsh.model.occ.addLine(p1, p2)  # Bottom edge
    l2 = gmsh.model.occ.addLine(p2, p3)  # Right edge
    l3 = gmsh.model.occ.addLine(p3, p4)  # Top edge
    l4 = gmsh.model.occ.addLine(p4, p1)  # Left edge# Define a rectangle in the XY plane

    loop = gmsh.model.occ.addCurveLoop([l1, l2, l3, l4])

    return loop, p1, p2, p3, p4, l1, l2, l3, l4

#%%
import gmsh
# Initialize GMSH
gmsh.initialize()

#%% 

""" INPUT dots in structure """

N_dots = 1
Names = ["Dot_1","Dot_2"]

input_dict = {
    Names[0]: {
        "shape": "circle",
        "x": 0,
        "y": 0.8,
        "radius": 0.05  # For cylinder
    },
    Names[1]: {
        "shape": "rectangle",
        "x": -0.5,
        "y": -0.25,
        "width": 1,
        "height": 0.5
    }
    # Add more dots as needed
}

#%% 
""" Plasma volume geometry """
# Cartesian coordinates of bottom plasma volume surface lower left corner
x_plasma_volume_ll = -2
y_plasma_volume_ll = -2
z_plasma_volume_ll = 0

# Width and height of plasma volume base surface 
width = 4
height = 4
dz_plasma_volume = 4  # This variable is defined but not used in this surface creation

# Create a curve loop for plasma volume base
plasma_base_rectangle_loop, p1, p2, p3, p4, l1, l2, l3, l4 = \
    rectangle_def(x_plasma_volume_ll, y_plasma_volume_ll, z_plasma_volume_ll, width, height) 

# Create a curve loop for plasma_volume top

plasma_top_rectangle_loop, p5, p6, p7, p8, l5, l6, l7, l8 = rectangle_def(x_plasma_volume_ll, y_plasma_volume_ll, \
                                                          z_plasma_volume_ll + dz_plasma_volume, width, height) 

# Create vertical lines connecting bottom and top
l9 = gmsh.model.occ.addLine(p1, p5)
l10 = gmsh.model.occ.addLine(p2, p6)
l11 = gmsh.model.occ.addLine(p3, p7)
l12 = gmsh.model.occ.addLine(p4, p8)

# list of surfaces IDs enclosing the plasma volume
volumes_surfaces = []

# Create surfaces for the sides (lateral surfaces)
plasma_side1 = gmsh.model.occ.addPlaneSurface([gmsh.model.occ.addCurveLoop([l1, l10, -l5, -l9])]) # minus when instead of going for extremity A to B you go from B to A close a loop
plasma_side2 = gmsh.model.occ.addPlaneSurface([gmsh.model.occ.addCurveLoop([l2, l11, -l6, -l10])])
plasma_side3 = gmsh.model.occ.addPlaneSurface([gmsh.model.occ.addCurveLoop([l3, l12, -l7, -l11])])
plasma_side4 = gmsh.model.occ.addPlaneSurface([gmsh.model.occ.addCurveLoop([l4, l9, -l8, -l12])])

# store surfaces enclosing volume in a variable
volumes_surfaces.append(plasma_side1)
volumes_surfaces.append(plasma_side2)
volumes_surfaces.append(plasma_side3)
volumes_surfaces.append(plasma_side4)

# Create the top surface
plasma_top = gmsh.model.occ.addPlaneSurface([gmsh.model.occ.addCurveLoop([l5, l6, l7, l8])])

# store surfaces enclosing volume in a variable
volumes_surfaces.append(plasma_top)

# Synchronize the GMSH model
gmsh.model.occ.synchronize()

#%% 
""" DiMES geometry """

# Cartesian coordinates of the center of the DiMES disk
x_dimes = 0
y_dimes = 0
z_dimes = 1  

# Radius of DiMES
r_dimes = 1

DiMES_circle = gmsh.model.occ.addCircle(x_dimes , y_dimes , 0 , r_dimes)

DiMES_circle_loop = gmsh.model.occ.addCurveLoop([DiMES_circle])

plasma_base = gmsh.model.occ.addPlaneSurface([plasma_base_rectangle_loop, DiMES_circle_loop]) 

# store surfaces enclosing volume in a variable
volumes_surfaces.append(plasma_base)

if z_dimes > 0:                  

    DiMES_edge_surface = gmsh.model.occ.extrude([(1,DiMES_circle)], 0, 0, z_dimes)

    # store surfaces enclosing volume in a variable
    for tup in DiMES_edge_surface:
        if tup[0] == 2:
            DiMES_edge_surface_id = tup[1]
            break

    volumes_surfaces.append(DiMES_edge_surface_id)

    #identify top_circle loop
    DiMES_top_circle_loop = gmsh.model.occ.addCurveLoop([DiMES_edge_surface[0][1]])

else:
    DiMES_top_circle_loop = DiMES_circle_loop

#%%

""" geometry Dots (coatings)"""

dot_loops = []

for nm in Names:

    if input_dict[nm]["shape"] == "circle":

        x = input_dict[nm]['x']
        y = input_dict[nm]['y']
        z = z_dimes
        r = input_dict[nm]['radius']

        dot_shape = gmsh.model.occ.addCircle(x, y, z, r)

        dot_loop = gmsh.model.occ.addCurveLoop([dot_shape])

    elif input_dict[nm]["shape"] == "rectangle":

        x = input_dict[nm]['x']
        y = input_dict[nm]['y']
        z = z_dimes
        width = input_dict[nm]['width']
        height = input_dict[nm]['height']

        dot_loop = rectangle_def(x, y, z, width, height)
        dot_loop = dot_loop[0]

    # store dots loops IDs
    dot_loops.append(dot_loop)
    # create dots surfaces
    gmsh.model.occ.addPlaneSurface([dot_loop])

    # store surfaces enclosing volume in a variable
    volumes_surfaces.append(gmsh.model.occ.addPlaneSurface([dot_loop]))

# Generate DiMES top surface
gmsh.model.occ.synchronize()
DiMES_top_surface = gmsh.model.occ.addPlaneSurface([DiMES_top_circle_loop] + dot_loops)

# store surfaces enclosing volume in a variable
volumes_surfaces.append(DiMES_top_surface)

#%% generate the volume
plasma_volume = gmsh.model.occ.addVolume([gmsh.model.occ.addSurfaceLoop(volumes_surfaces)])

#%% Generate the mesh and visualize the result

# Final synchronization of the CAD model
gmsh.model.occ.synchronize()

gmsh.option.setNumber("Mesh.MeshSizeFromCurvature", 1)
gmsh.option.setNumber("Mesh.MinimumElementsPerTwoPi", 20)

# # Prevent very small elements in small dots
gmsh.option.setNumber("Mesh.MeshSizeMin", 0.05)
# Set maximum mesh characteristic length for the whole model
gmsh.option.setNumber("Mesh.CharacteristicLengthMax", 0.2) 
gmsh.model.mesh.generate(2)

# Launch the GUI to see the results:
# Optionally, run the GUI to visualize
gmsh.fltk.run()

# Finalize GMSH
gmsh.finalize()
LC-GA commented 1 month ago

this code allows you to simulate DiMES in a plasma volume, only rectangles and circles are allowed as dots geometries. DiMES z location can be manually adjusted. Warning: the code doesn't tell you if your input is correct. Pay attention before simulating your geometry