diff --git a/UM2N/generator/__init__.py b/UM2N/generator/__init__.py
index fcf7d7e..bf4132f 100644
--- a/UM2N/generator/__init__.py
+++ b/UM2N/generator/__init__.py
@@ -3,7 +3,6 @@
 from .rand_source_generator import *  # noqa
 from .equation_generator import *  # noqa
 
-from .polymesh import *  # noqa
-from .squaremesh import *  # noqa
+from .unstructured_mesh import *  # noqa
 from .burgers_solver import *  # noqa
 from .swirl_solver import *  # noqa
diff --git a/UM2N/generator/polymesh.py b/UM2N/generator/polymesh.py
deleted file mode 100644
index e45ecb4..0000000
--- a/UM2N/generator/polymesh.py
+++ /dev/null
@@ -1,154 +0,0 @@
-import random
-
-import firedrake as fd
-import gmsh
-import numpy as np
-
-__all__ = ["RandPolyMesh"]
-
-
-class RandPolyMesh:
-    """
-    Create a random polygonal mesh by spliting the edge of a
-    square randomly.
-    """
-
-    def __init__(self, scale=1.0, mesh_type=2):
-        # params setup
-        self.mesh_type = mesh_type
-        self.scale = scale
-        self.start = 0
-        self.end = self.scale
-        self.split_threshold = 0.3
-        self.mid = (self.start + self.end) / 2
-        self.quater = (self.start + self.mid) / 2
-        self.three_quater = (self.mid + self.end) / 2
-        self.mid_interval = (self.end - self.start) / 3
-        self.quater_interval = (self.mid - self.start) / 4
-        # temp vars
-        self.points = []
-        self.lines = []
-        # generate mesh
-        self.get_rand_points()
-        return
-
-    def get_mesh(self, res=1e-1, file_path="./temp.msh"):
-        gmsh.initialize()
-        gmsh.model.add("t1")
-        # params setup
-        self.lc = res
-        self.start = 0
-        self.end = self.scale
-        self.mid = (self.start + self.end) / 2
-        self.quater = (self.start + self.mid) / 2
-        self.three_quater = (self.mid + self.end) / 2
-        self.mid_interval = (self.end - self.start) / 3
-        self.quater_interval = (self.mid - self.start) / 4
-        self.file_path = file_path
-        # temp vars
-        self.points = []
-        self.lines = []
-        # generate mesh
-        self.get_points()
-        self.get_line()
-        self.get_curve()
-        self.get_plane()
-        gmsh.model.geo.synchronize()
-        gmsh.option.setNumber("Mesh.Algorithm", self.mesh_type)
-        self.get_boundaries()
-        gmsh.model.addPhysicalGroup(2, [1], name="My surface")
-        gmsh.model.mesh.generate(2)
-        gmsh.write(self.file_path)
-        gmsh.finalize()
-        self.num_boundary = len(self.lines)
-        return fd.Mesh(self.file_path)
-
-    def get_rand(self, mean, interval):
-        return random.uniform(mean - interval, mean + interval)
-
-    def get_rand_points(self):
-        points = []
-        split_p = np.random.uniform(0, 1, 4)
-        # edge 1
-        if split_p[0] < self.split_threshold:
-            points.append([self.get_rand(self.quater, self.quater_interval), 0])
-            points.append([self.get_rand(self.three_quater, self.quater_interval), 0])
-        else:
-            points.append([self.get_rand(self.mid, self.mid_interval), 0])
-        # edge 2
-        if split_p[1] < self.split_threshold:
-            points.append(
-                [self.scale, self.get_rand(self.quater, self.quater_interval)]
-            )
-            points.append(
-                [self.scale, self.get_rand(self.three_quater, self.quater_interval)]
-            )
-        else:
-            points.append([self.scale, self.get_rand(self.mid, self.mid_interval)])
-        # edge 3
-        if split_p[2] < self.split_threshold:
-            points.append(
-                [self.get_rand(self.three_quater, self.quater_interval), self.scale]
-            )
-            points.append(
-                [self.get_rand(self.quater, self.quater_interval), self.scale]
-            )
-        else:
-            points.append([self.get_rand(self.mid, self.mid_interval), self.scale])
-        # edge 4
-        if split_p[3] < self.split_threshold:
-            points.append([0, self.get_rand(self.three_quater, self.quater_interval)])
-            points.append([0, self.get_rand(self.quater, self.quater_interval)])
-        else:
-            points.append([0, self.get_rand(self.mid, self.mid_interval)])
-            # points.append(p1)
-        self.raw_points = points
-        return
-
-    def get_points(self):
-        temp = []
-        for i in range(len(self.raw_points)):
-            temp.append(
-                gmsh.model.geo.addPoint(
-                    self.raw_points[i][0], self.raw_points[i][1], 0, self.lc
-                )
-            )
-        self.points = temp
-
-    def get_line(self):
-        for i in range(len(self.points)):
-            if i < len(self.points) - 1:
-                line = gmsh.model.geo.addLine(self.points[i], self.points[i + 1])
-                self.lines.append(line)
-            else:
-                line = gmsh.model.geo.addLine(self.points[i], self.points[0])
-                self.lines.append(line)
-        return
-
-    def get_boundaries(self):
-        print("in get_boundaries lines:", self.lines)
-        for i, line_tag in enumerate(self.lines):
-            gmsh.model.addPhysicalGroup(1, [line_tag], i + 1)
-            gmsh.model.setPhysicalName(1, i + 1, "Boundary " + str(i + 1))
-
-    def get_curve(self):
-        gmsh.model.geo.addCurveLoop([i for i in range(1, len(self.points) + 1)], 1)
-
-    def get_plane(self):
-        gmsh.model.geo.addPlaneSurface([1], 1)
-
-    def show(self, file_path):
-        mesh = fd.Mesh(file_path)
-        fig = fd.triplot(mesh)
-        return fig
-
-
-if __name__ == "__main__":
-    import matplotlib.pyplot as plt
-
-    mesh_gen = RandPolyMesh(mesh_type=2)
-    mesh_coarse = mesh_gen.get_mesh(res=5e-2, file_path="./temp1.msh")
-    mesh_fine = mesh_gen.get_mesh(res=4e-2, file_path="./temp2.msh")
-    mesh_gen.show("./temp1.msh")
-    mesh_gen.show("./temp2.msh")
-    plt.show()
diff --git a/UM2N/generator/squaremesh.py b/UM2N/generator/squaremesh.py
deleted file mode 100644
index 98254eb..0000000
--- a/UM2N/generator/squaremesh.py
+++ /dev/null
@@ -1,109 +0,0 @@
-import firedrake as fd
-import gmsh
-
-__all__ = ["UnstructuredSquareMesh"]
-
-
-class UnstructuredSquareMesh:
-    """
-    Create a random polygonal mesh by spliting the edge of a
-    square randomly.
-    """
-
-    def __init__(self, scale=1.0, mesh_type=2):
-        # params setup
-        self.mesh_type = mesh_type
-        self.scale = scale
-        self.start = 0
-        self.end = self.scale
-
-        self.points = []
-        self.lines = []
-        return
-
-    def get_mesh(self, res=1e-1, file_path="./temp.msh"):
-        gmsh.initialize()
-        gmsh.model.add("t1")
-        # params setup
-        self.lc = res
-        self.start = 0
-        self.end = self.scale
-        self.file_path = file_path
-        # temp vars
-        self.points = []
-        self.lines = []
-        # generate mesh
-        self.get_corner_points()
-        self.get_points()
-        self.get_line()
-        self.get_curve()
-        self.get_plane()
-        gmsh.model.geo.synchronize()
-        gmsh.option.setNumber("Mesh.Algorithm", self.mesh_type)
-        self.get_boundaries()
-        gmsh.model.addPhysicalGroup(2, [1], name="My surface")
-        gmsh.model.mesh.generate(2)
-        gmsh.write(self.file_path)
-        gmsh.finalize()
-        self.num_boundary = len(self.lines)
-        return fd.Mesh(self.file_path)
-
-    def get_corner_points(self):
-        points = []
-        points.append([0, 0])
-        points.append([1, 0])
-        points.append([1, 1])
-        points.append([0, 1])
-        self.raw_points = points
-        return
-
-    def get_points(self):
-        temp = []
-        for i in range(len(self.raw_points)):
-            temp.append(
-                gmsh.model.geo.addPoint(
-                    self.raw_points[i][0], self.raw_points[i][1], 0, self.lc
-                )
-            )
-        self.points = temp
-
-    def get_line(self):
-        for i in range(len(self.points)):
-            if i < len(self.points) - 1:
-                line = gmsh.model.geo.addLine(self.points[i], self.points[i + 1])
-                self.lines.append(line)
-            else:
-                line = gmsh.model.geo.addLine(self.points[i], self.points[0])
-                self.lines.append(line)
-        return
-
-    def get_boundaries(self):
-        print("in get_boundaries lines:", self.lines)
-        for i, line_tag in enumerate(self.lines):
-            gmsh.model.addPhysicalGroup(1, [line_tag], i + 1)
-            gmsh.model.setPhysicalName(1, i + 1, "Boundary " + str(i + 1))
-
-    def get_curve(self):
-        gmsh.model.geo.addCurveLoop([i for i in range(1, len(self.points) + 1)], 1)
-
-    def get_plane(self):
-        gmsh.model.geo.addPlaneSurface([1], 1)
-
-    def show(self, file_path):
-        mesh = fd.Mesh(file_path)
-        fig = fd.triplot(mesh)
-        return fig
-
-    def load_mesh(self, file_path):
-        return fd.Mesh(file_path)
-
-
-if __name__ == "__main__":
-    import matplotlib.pyplot as plt
-
-    mesh_gen = UnstructuredSquareMesh(mesh_type=1)
-    mesh_coarse = mesh_gen.get_mesh(res=5e-2, file_path="./temp1.msh")
-    mesh_fine = mesh_gen.get_mesh(res=4e-2, file_path="./temp2.msh")
-    mesh_gen.show("./temp1.msh")
-    mesh_gen.show("./temp2.msh")
-    plt.show()
diff --git a/UM2N/generator/unstructured_mesh.py b/UM2N/generator/unstructured_mesh.py
new file mode 100644
index 0000000..d617116
--- /dev/null
+++ b/UM2N/generator/unstructured_mesh.py
@@ -0,0 +1,196 @@
+"""
+Module for handling generating unstructured meshes.
+"""
+
+import abc
+import os
+import random
+
+import gmsh
+import numpy as np
+from firedrake.mesh import Mesh
+
+__all__ = [
+    "UnstructuredSquareMeshGenerator",
+    "UnstructuredRandomPolygonalMeshGenerator",
+]
+
+
+class UnstructuredMeshGenerator(abc.ABC):
+    """
+    Base class for mesh generators.
+    """
+
+    def __init__(self, scale=1.0, mesh_type=2):
+        """
+        :kwarg scale: overall scale factor for the domain size (default: 1.0)
+        :type scale: float
+        :kwarg mesh_type: Gmsh algorithm number (default: 2)
+        :type mesh_type: int
+        """
+        self.scale = scale
+        # TODO: More detail on Gmsh algorithm number (#50)
+        self.mesh_type = mesh_type
+        self._mesh = None
+
+    @property
+    @abc.abstractmethod
+    def corners(self):
+        """
+        Property defining the coordinates of the corner vertices of the domain to be
+        meshed.
+
+        :returns: coordinates of the corner vertices of the domain
+        :rtype: tuple
+        """
+        pass
+
+    def generate_mesh(self, res=0.1, output_filename="./temp.msh", remove_file=False):
+        """
+        Generate a mesh at a given resolution level.
+
+        :kwarg res: mesh resolution (element diameter) (default: 0.1, suitable for mesh
+            with scale 1.0)
+        :type res: float
+        :kwarg output_filename: filename for saving the mesh, including the path and .msh
+            extension (default: './temp.msh')
+        :type output_filename: str
+        :kwarg remove_file: should the .msh file be removed after generation? (default:
+            False)
+        :type remove_file: bool
+        :returns: mesh generated
+        :rtype: :class:`firedrake.mesh.MeshGeometry`
+        """
+        gmsh.initialize()
+        gmsh.model.add("t1")
+        self.lc = res
+        self._points = [
+            gmsh.model.geo.addPoint(*corner, 0, self.lc) for corner in self.corners
+        ]
+        self._lines = [
+            gmsh.model.geo.addLine(point, point_next)
+            for point, point_next in zip(
+                self._points, self._points[1:] + [self._points[0]]
+            )
+        ]
+        gmsh.model.geo.addCurveLoop([i + 1 for i in range(len(self._points))], 1)
+        gmsh.model.geo.addPlaneSurface([1], 1)
+        gmsh.model.geo.synchronize()
+        gmsh.option.setNumber("Mesh.Algorithm", self.mesh_type)
+        for i, line_tag in enumerate(self._lines):
+            gmsh.model.addPhysicalGroup(1, [line_tag], i + 1)
+            gmsh.model.setPhysicalName(1, i + 1, "Boundary " + str(i + 1))
+        gmsh.model.addPhysicalGroup(2, [1], name="My surface")
+        gmsh.model.mesh.generate(2)
+        gmsh.write(output_filename)
+        gmsh.finalize()
+        self.num_boundary = len(self._lines)
+        self._mesh = Mesh(output_filename)
+        if remove_file:
+            os.remove(output_filename)
+        return self._mesh
+
+    def load_mesh(self, filename):
+        """
+        Load a mesh from a file saved in .msh format.
+
+        :arg filename: filename including path and the .msh extension
+        :type filename: str
+        :returns: mesh loaded from file
+        :rtype: :class:`firedrake.mesh.MeshGeometry`
+        """
+        self._mesh = Mesh(filename)
+        return self._mesh
+
+
+class UnstructuredSquareMeshGenerator(UnstructuredMeshGenerator):
+    """
+    Generate an unstructured mesh of a 2D square domain using Gmsh.
+    """
+
+    @property
+    def corners(self):
+        """
+        Property defining the coordinates of the corner vertices of the domain to be
+        meshed.
+
+        :returns: coordinates of the corner vertices of the domain
+        :rtype: tuple
+        """
+        return ((0, 0), (self.scale, 0), (self.scale, self.scale), (0, self.scale))
+
+
+class UnstructuredRandomPolygonalMeshGenerator(UnstructuredMeshGenerator):
+    """
+    Create a random polygonal mesh by spliting the edge of a
+    square randomly.
+    """
+
+    @staticmethod
+    def sample_uniform(mean, interval):
+        """
+        Sample a point from a uniform distribution with a given mean and interval.
+
+        Note that the interval is *either side* of the mean, not the overall interval,
+        i.e., overall interval is symmetric: `[mean - interval, mean + interval]`.
+
+        :arg mean: the mean value of the uniform distribution
+        :type mean: float
+        :arg interval: the interval either side of the mean
+        :type interval: float
+        :returns: sampled point
+        :rtype: float
+        """
+        return random.uniform(mean - interval, mean + interval)
+
+    @property
+    def corners(self):
+        """
+        Property defining the coordinates of the corner vertices of the domain to be
+        meshed.
+
+        :returns: coordinates of the corner vertices of the domain
+        :rtype: tuple
+        """
+        if hasattr(self, "_corners"):
+            return self._corners
+        start = 0
+        finish = self.scale
+        split_threshold = 0.3
+        mid = (start + finish) / 2
+        quarter = (start + mid) / 2
+        three_quarter = (mid + finish) / 2
+        mid_interval = (finish - start) / 3
+        quarter_interval = (mid - start) / 4
+        points = []
+        split_p = np.random.uniform(0, 1, 4)
+        # edge 1
+        if split_p[0] < split_threshold:
+            points.append([self.sample_uniform(quarter, quarter_interval), 0])
+            points.append([self.sample_uniform(three_quarter, quarter_interval), 0])
+        else:
+            points.append([self.sample_uniform(mid, mid_interval), 0])
+        # edge 2
+        if split_p[1] < split_threshold:
+            points.append([self.scale, self.sample_uniform(quarter, quarter_interval)])
+            points.append(
+                [self.scale, self.sample_uniform(three_quarter, quarter_interval)]
+            )
+        else:
+            points.append([self.scale, self.sample_uniform(mid, mid_interval)])
+        # edge 3
+        if split_p[2] < split_threshold:
+            points.append(
+                [self.sample_uniform(three_quarter, quarter_interval), self.scale]
+            )
+            points.append([self.sample_uniform(quarter, quarter_interval), self.scale])
+        else:
+            points.append([self.sample_uniform(mid, mid_interval), self.scale])
+        # edge 4
+        if split_p[3] < split_threshold:
+            points.append([0, self.sample_uniform(three_quarter, quarter_interval)])
+            points.append([0, self.sample_uniform(quarter, quarter_interval)])
+        else:
+            points.append([0, self.sample_uniform(mid, mid_interval)])
+        self._corners = tuple(points)
+        return self._corners
diff --git a/UM2N/model/train_util.py b/UM2N/model/train_util.py
index 2404fc5..3ade909 100644
--- a/UM2N/model/train_util.py
+++ b/UM2N/model/train_util.py
@@ -5,7 +5,18 @@
 import numpy as np
 import torch
 import torch.nn as nn
-from pytorch3d.loss import chamfer_distance
+
+# FIXME: Hack to handle the case where pytorch3d cannot be installed (#41)
+try:
+    from pytorch3d.loss import chamfer_distance
+except ImportError:
+
+    def chamfer_distance(*args):
+        raise ImportError(
+            "chamfer_distance depends on pytorch3d, which cannot be imported"
+        )
+
+
 from torch_geometric.loader import DataLoader
 from torch_geometric.nn import MessagePassing, knn_graph
 
diff --git a/play/play_grad.ipynb b/play/play_grad.ipynb
index 919ac90..94c4bf2 100644
--- a/play/play_grad.ipynb
+++ b/play/play_grad.ipynb
@@ -6,13 +6,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import torch\n",
+    "import warnings\n",
+    "\n",
     "import firedrake as fd\n",
-    "import warpmesh as wm\n",
     "import movement as mv\n",
+    "import torch\n",
+    "import warpmesh as wm\n",
     "from torch_geometric.data import DataLoader\n",
     "\n",
-    "import warnings\n",
     "warnings.filterwarnings('ignore')\n",
     "device = torch.device('cuda' if torch.cuda.is_available()\n",
     "else 'cpu')"
diff --git a/pyproject.toml b/pyproject.toml
index 8171dd1..28b35fc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -59,5 +59,4 @@ ignore = [
   "E402",  # module level import not at top of file
   "F403",  # unable to detect undefined names
   "F405",  # name may be undefined, or defined from star imports
-  "I001",  # import block is unsorted or unformatted
 ]
diff --git a/script/build_burgers_square.py b/script/build_burgers_square.py
index 287ae06..a8ed4ff 100644
--- a/script/build_burgers_square.py
+++ b/script/build_burgers_square.py
@@ -402,15 +402,16 @@ def sample_from_loop(
             mesh_new = None
             mesh_fine = None
             if mesh_type != 0:
-                unstructure_square_mesh_gen = UM2N.UnstructuredSquareMesh(
+                unstructured_square_mesh_gen = UM2N.UnstructuredSquareMesh(
                     scale=scale_x, mesh_type=mesh_type
                 )  # noqa
-                mesh = unstructure_square_mesh_gen.get_mesh(
-                    res=lc, file_path=os.path.join(problem_mesh_dir, "mesh.msh")
+                mesh = unstructured_square_mesh_gen.generate_mesh(
+                    res=lc, output_filename=os.path.join(problem_mesh_dir, "mesh.msh")
                 )
                 mesh_new = fd.Mesh(os.path.join(problem_mesh_dir, "mesh.msh"))
-                mesh_fine = unstructure_square_mesh_gen.get_mesh(
-                    res=1e-2, file_path=os.path.join(problem_mesh_fine_dir, "mesh.msh")
+                mesh_fine = unstructured_square_mesh_gen.generate_mesh(
+                    res=1e-2,
+                    output_filename=os.path.join(problem_mesh_fine_dir, "mesh.msh"),
                 )
             else:
                 mesh = fd.UnitSquareMesh(n_grid, n_grid)
diff --git a/script/build_helmholtz_poly.py b/script/build_helmholtz_poly.py
index 1d11c87..080116f 100644
--- a/script/build_helmholtz_poly.py
+++ b/script/build_helmholtz_poly.py
@@ -216,9 +216,11 @@ def move_data(target, source, start, num_file):
     while i < n_samples:
         try:
             print("Generating Sample: " + str(i))
-            rand_poly_mesh_gen = UM2N.RandPolyMesh(scale=scale_x, mesh_type=mesh_type)  # noqa
-            mesh = rand_poly_mesh_gen.get_mesh(
-                res=lc, file_path=os.path.join(problem_mesh_dir, f"mesh{i}.msh")
+            rand_poly_mesh_gen = UM2N.UnstructuredRandomPolygonalMeshGenerator(
+                scale=scale_x, mesh_type=mesh_type
+            )  # noqa
+            mesh = rand_poly_mesh_gen.generate_mesh(
+                res=lc, output_filename=os.path.join(problem_mesh_dir, f"mesh{i}.msh")
             )
             num_boundary = rand_poly_mesh_gen.num_boundary
             # Generate Random solution field
@@ -258,8 +260,9 @@ def move_data(target, source, start, num_file):
             hessian = UM2N.MeshGenerator(
                 params={
                     "eq": helmholtz_eq,
-                    "mesh": rand_poly_mesh_gen.get_mesh(
-                        res=lc, file_path=os.path.join(problem_mesh_dir, f"mesh{i}.msh")
+                    "mesh": rand_poly_mesh_gen.generate_mesh(
+                        res=lc,
+                        output_filename=os.path.join(problem_mesh_dir, f"mesh{i}.msh"),
                     ),
                 }
             ).get_hessian(mesh)
@@ -267,8 +270,9 @@ def move_data(target, source, start, num_file):
             hessian_norm = UM2N.MeshGenerator(
                 params={
                     "eq": helmholtz_eq,
-                    "mesh": rand_poly_mesh_gen.get_mesh(
-                        res=lc, file_path=os.path.join(problem_mesh_dir, f"mesh{i}.msh")
+                    "mesh": rand_poly_mesh_gen.generate_mesh(
+                        res=lc,
+                        output_filename=os.path.join(problem_mesh_dir, f"mesh{i}.msh"),
                     ),
                 }
             ).monitor_func(mesh)
@@ -281,8 +285,9 @@ def move_data(target, source, start, num_file):
             mesh_gen = UM2N.MeshGenerator(
                 params={
                     "eq": helmholtz_eq,
-                    "mesh": rand_poly_mesh_gen.get_mesh(
-                        res=lc, file_path=os.path.join(problem_mesh_dir, f"mesh{i}.msh")
+                    "mesh": rand_poly_mesh_gen.generate_mesh(
+                        res=lc,
+                        output_filename=os.path.join(problem_mesh_dir, f"mesh{i}.msh"),
                     ),
                 }
             )
@@ -381,8 +386,9 @@ def move_data(target, source, start, num_file):
             # ==========================================
 
             # generate log file
-            high_res_mesh = rand_poly_mesh_gen.get_mesh(
-                res=1e-2, file_path=os.path.join(problem_mesh_fine_dir, f"mesh{i}.msh")
+            high_res_mesh = rand_poly_mesh_gen.generate_mesh(
+                res=1e-2,
+                output_filename=os.path.join(problem_mesh_fine_dir, f"mesh{i}.msh"),
             )
 
             high_res_function_space = fd.FunctionSpace(high_res_mesh, "CG", 1)
diff --git a/script/build_helmholtz_square.py b/script/build_helmholtz_square.py
index 447368d..c8d25d1 100644
--- a/script/build_helmholtz_square.py
+++ b/script/build_helmholtz_square.py
@@ -234,12 +234,12 @@ def move_data(target, source, start, num_file):
         try:
             print("Generating Sample: " + str(i))
             if mesh_type != 0:
-                unstructure_square_mesh_gen = UM2N.UnstructuredSquareMesh(
+                unstructured_square_mesh_gen = UM2N.UnstructuredSquareMesh(
                     scale=scale_x, mesh_type=mesh_type
                 )  # noqa
-                mesh = unstructure_square_mesh_gen.get_mesh(
+                mesh = unstructured_square_mesh_gen.generate_mesh(
                     res=lc,
-                    file_path=os.path.join(problem_mesh_dir, f"mesh_{i:04d}.msh"),
+                    output_filename=os.path.join(problem_mesh_dir, f"mesh_{i:04d}.msh"),
                 )
             else:
                 n_grid = int(1 / lc)
@@ -427,9 +427,11 @@ def move_data(target, source, start, num_file):
 
             if mesh_type != 0:
                 # generate log file
-                high_res_mesh = unstructure_square_mesh_gen.get_mesh(
+                high_res_mesh = unstructured_square_mesh_gen.generate_mesh(
                     res=1e-2,
-                    file_path=os.path.join(problem_mesh_fine_dir, f"mesh_{i:04d}.msh"),
+                    output_filename=os.path.join(
+                        problem_mesh_fine_dir, f"mesh_{i:04d}.msh"
+                    ),
                 )
             else:
                 high_res_mesh = fd.UnitSquareMesh(100, 100)
diff --git a/script/build_poisson_poly.py b/script/build_poisson_poly.py
index 03a0fcf..87634f6 100644
--- a/script/build_poisson_poly.py
+++ b/script/build_poisson_poly.py
@@ -216,9 +216,11 @@ def move_data(target, source, start, num_file):
     while i < n_samples:
         try:
             print("Generating Sample: " + str(i))
-            rand_poly_mesh_gen = UM2N.RandPolyMesh(scale=scale_x, mesh_type=mesh_type)  # noqa
-            mesh = rand_poly_mesh_gen.get_mesh(
-                res=lc, file_path=os.path.join(problem_mesh_dir, f"mesh{i}.msh")
+            rand_poly_mesh_gen = UM2N.UnstructuredRandomPolygonalMeshGenerator(
+                scale=scale_x, mesh_type=mesh_type
+            )  # noqa
+            mesh = rand_poly_mesh_gen.generate_mesh(
+                res=lc, output_filename=os.path.join(problem_mesh_dir, f"mesh{i}.msh")
             )
             num_boundary = rand_poly_mesh_gen.num_boundary
             # Generate Random solution field
@@ -256,8 +258,9 @@ def move_data(target, source, start, num_file):
             hessian = UM2N.MeshGenerator(
                 params={
                     "eq": poisson_eq,
-                    "mesh": rand_poly_mesh_gen.get_mesh(
-                        res=lc, file_path=os.path.join(problem_mesh_dir, f"mesh{i}.msh")
+                    "mesh": rand_poly_mesh_gen.generate_mesh(
+                        res=lc,
+                        output_filename=os.path.join(problem_mesh_dir, f"mesh{i}.msh"),
                     ),
                 }
             ).get_hessian(mesh)
@@ -265,8 +268,9 @@ def move_data(target, source, start, num_file):
             hessian_norm = UM2N.MeshGenerator(
                 params={
                     "eq": poisson_eq,
-                    "mesh": rand_poly_mesh_gen.get_mesh(
-                        res=lc, file_path=os.path.join(problem_mesh_dir, f"mesh{i}.msh")
+                    "mesh": rand_poly_mesh_gen.generate_mesh(
+                        res=lc,
+                        output_filename=os.path.join(problem_mesh_dir, f"mesh{i}.msh"),
                     ),
                 }
             ).monitor_func(mesh)
@@ -279,8 +283,9 @@ def move_data(target, source, start, num_file):
             mesh_gen = UM2N.MeshGenerator(
                 params={
                     "eq": poisson_eq,
-                    "mesh": rand_poly_mesh_gen.get_mesh(
-                        res=lc, file_path=os.path.join(problem_mesh_dir, f"mesh{i}.msh")
+                    "mesh": rand_poly_mesh_gen.generate_mesh(
+                        res=lc,
+                        output_filename=os.path.join(problem_mesh_dir, f"mesh{i}.msh"),
                     ),
                 }
             )
@@ -378,8 +383,9 @@ def move_data(target, source, start, num_file):
             # ==========================================
 
             # generate log file
-            high_res_mesh = rand_poly_mesh_gen.get_mesh(
-                res=1e-2, file_path=os.path.join(problem_mesh_fine_dir, f"mesh{i}.msh")
+            high_res_mesh = rand_poly_mesh_gen.generate_mesh(
+                res=1e-2,
+                output_filename=os.path.join(problem_mesh_fine_dir, f"mesh{i}.msh"),
             )
 
             high_res_function_space = fd.FunctionSpace(high_res_mesh, "CG", 1)
diff --git a/script/build_poisson_square.py b/script/build_poisson_square.py
index 0051279..4bf9712 100644
--- a/script/build_poisson_square.py
+++ b/script/build_poisson_square.py
@@ -214,11 +214,11 @@ def move_data(target, source, start, num_file):
     while i < n_samples:
         try:
             print("Generating Sample: " + str(i))
-            unstructure_square_mesh_gen = UM2N.UnstructuredSquareMesh(
+            unstructured_square_mesh_gen = UM2N.UnstructuredSquareMesh(
                 scale=scale_x, mesh_type=mesh_type
             )  # noqa
-            mesh = unstructure_square_mesh_gen.get_mesh(
-                res=lc, file_path=os.path.join(problem_mesh_dir, f"mesh{i}.msh")
+            mesh = unstructured_square_mesh_gen.generate_mesh(
+                res=lc, output_filename=os.path.join(problem_mesh_dir, f"mesh{i}.msh")
             )
             # Generate Random solution field
             rand_u_generator = UM2N.RandSourceGenerator(
@@ -374,8 +374,9 @@ def move_data(target, source, start, num_file):
             # ==========================================
 
             # generate log file
-            high_res_mesh = unstructure_square_mesh_gen.get_mesh(
-                res=1e-2, file_path=os.path.join(problem_mesh_fine_dir, f"mesh{i}.msh")
+            high_res_mesh = unstructured_square_mesh_gen.generate_mesh(
+                res=1e-2,
+                output_filename=os.path.join(problem_mesh_fine_dir, f"mesh{i}.msh"),
             )
             high_res_function_space = fd.FunctionSpace(high_res_mesh, "CG", 1)
 
diff --git a/script/build_swirl.py b/script/build_swirl.py
index 3c9ebf1..c6e0859 100644
--- a/script/build_swirl.py
+++ b/script/build_swirl.py
@@ -419,18 +419,18 @@ def sample_from_loop(
     mesh_new = None
     if mesh_type != 0:
         mesh_gen = UM2N.UnstructuredSquareMesh(mesh_type=mesh_type)
-        mesh = mesh_gen.get_mesh(
-            res=lc, file_path=os.path.join(problem_mesh_dir, "mesh.msh")
+        mesh = mesh_gen.generate_mesh(
+            res=lc, output_filename=os.path.join(problem_mesh_dir, "mesh.msh")
         )
-        mesh_new = mesh_gen.get_mesh(
-            res=lc, file_path=os.path.join(problem_mesh_dir, "mesh.msh")
+        mesh_new = mesh_gen.generate_mesh(
+            res=lc, output_filename=os.path.join(problem_mesh_dir, "mesh.msh")
         )
-        mesh_model = mesh_gen.get_mesh(
-            res=lc, file_path=os.path.join(problem_mesh_dir, "mesh.msh")
+        mesh_model = mesh_gen.generate_mesh(
+            res=lc, output_filename=os.path.join(problem_mesh_dir, "mesh.msh")
         )
         mesh_gen_fine = UM2N.UnstructuredSquareMesh(mesh_type=mesh_type)
-        mesh_fine = mesh_gen_fine.get_mesh(
-            res=1e-2, file_path=os.path.join(problem_mesh_fine_dir, "mesh.msh")
+        mesh_fine = mesh_gen_fine.generate_mesh(
+            res=1e-2, output_filename=os.path.join(problem_mesh_fine_dir, "mesh.msh")
         )
     else:
         mesh = fd.UnitSquareMesh(n_grid, n_grid)
diff --git a/script/sample_edges.ipynb b/script/sample_edges.ipynb
index 692cfe5..8ce2d95 100644
--- a/script/sample_edges.ipynb
+++ b/script/sample_edges.ipynb
@@ -72,13 +72,6 @@
     "                \"--add_nei=\" + str(add_nei),\n",
     "            ],check=True)\n"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
diff --git a/script/train_um2n.ipynb b/script/train_um2n.ipynb
index c1781f0..f885002 100644
--- a/script/train_um2n.ipynb
+++ b/script/train_um2n.ipynb
@@ -6,7 +6,7 @@
     "id": "egGPGDSxzRis"
    },
    "source": [
-    "# Training Notebok\n",
+    "# Training Notebook\n",
     "\n",
     "This notebook is aimed to utilise colab to train the model faster.\n",
     "Please upload this notebook to colab and to train the model.\n"
@@ -34,8 +34,9 @@
     "!wget -q https://raw.githubusercontent.com/chunyang-w/colab-github/main/colab_github.py\n",
     "\n",
     "import os\n",
-    "import torch\n",
+    "\n",
     "import colab_github\n",
+    "import torch\n",
     "\n",
     "os.environ['TORCH'] = torch.__version__\n",
     "print(torch.__version__)\n",
@@ -98,16 +99,17 @@
     "%load_ext autoreload\n",
     "%autoreload 1\n",
     "\n",
-    "from model import M2N, train, evaluate, MRN, count_dataset_tangle\n",
-    "from helper import mkdir_if_not_exist, plot_loss, plot_tangle\n",
-    "from loader import MeshDataset, normalise, AggreateDataset\n",
-    "from google.colab import runtime\n",
     "import os\n",
+    "import warnings\n",
+    "from datetime import datetime\n",
+    "\n",
+    "import pandas as pd\n",
     "import torch\n",
+    "from google.colab import runtime\n",
+    "from helper import mkdir_if_not_exist, plot_loss, plot_tangle\n",
+    "from loader import AggreateDataset, MeshDataset, normalise\n",
+    "from model import M2N, MRN, count_dataset_tangle, evaluate, train\n",
     "from torch.utils.tensorboard import SummaryWriter\n",
-    "import pandas as pd\n",
-    "from datetime import datetime\n",
-    "import warnings\n",
     "from torch_geometric.data import DataLoader\n",
     "\n",
     "warnings.filterwarnings(\"ignore\")\n",
@@ -636,15 +638,6 @@
    "source": [
     "runtime.unassign()"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "ixvccpsMPJtV"
-   },
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
diff --git a/temp.msh b/temp.msh
deleted file mode 100644
index d7043be..0000000
--- a/temp.msh
+++ /dev/null
@@ -1,529 +0,0 @@
-$MeshFormat
-4.1 0 8
-$EndMeshFormat
-$PhysicalNames
-7
-1 1 "Boundary 1"
-1 2 "Boundary 2"
-1 3 "Boundary 3"
-1 4 "Boundary 4"
-1 5 "Boundary 5"
-1 6 "Boundary 6"
-2 7 "My surface"
-$EndPhysicalNames
-$Entities
-6 6 1 0
-1 0.4634261501045485 0 0 0 
-2 1 0.1985221205847214 0 0 
-3 1 0.8523572185131376 0 0 
-4 0.7460765869539596 1 0 0 
-5 0.1467614907523769 1 0 0 
-6 0 0.6666386429395867 0 0 
-1 0.4634261501045485 0 0 1 0.1985221205847214 0 1 1 2 1 -2 
-2 1 0.1985221205847214 0 1 0.8523572185131376 0 1 2 2 2 -3 
-3 0.7460765869539596 0.8523572185131376 0 1 1 0 1 3 2 3 -4 
-4 0.1467614907523769 1 0 0.7460765869539596 1 0 1 4 2 4 -5 
-5 0 0.6666386429395867 0 0.1467614907523769 1 0 1 5 2 5 -6 
-6 0 0 0 0.4634261501045485 0.6666386429395867 0 1 6 2 6 -1 
-1 0 0 0 1 1 0 1 7 6 1 2 3 4 5 6 
-$EndEntities
-$Nodes
-13 119 1 119
-0 1 0 1
-1
-0.4634261501045485 0 0
-0 2 0 1
-2
-1 0.1985221205847214 0
-0 3 0 1
-3
-1 0.8523572185131376 0
-0 4 0 1
-4
-0.7460765869539596 1 0
-0 5 0 1
-5
-0.1467614907523769 1 0
-0 6 0 1
-6
-0 0.6666386429395867 0
-1 1 0 5
-7
-8
-9
-10
-11
-0.5528551250869682 0.03308702009739602 0
-0.6422841000692496 0.06617404019474085 0
-0.731713075051617 0.09926106029211751 0
-0.8211420500344074 0.1323480803896507 0
-0.9105710250171517 0.1654351004871668 0
-1 2 0 6
-12
-13
-14
-15
-16
-17
-1 0.2919271345743185 0
-1 0.3853321485638025 0
-1 0.4787371625532744 0
-1 0.5721421765430916 0
-1 0.6655471905331471 0
-1 0.7589522045231067 0
-1 3 0 2
-18
-19
-0.9153588623180809 0.9015714790087035 0
-0.8307177246361215 0.9507857395042929 0
-1 4 0 5
-20
-21
-22
-23
-24
-0.6461907375872316 1 0
-0.5463048882206047 1 0
-0.4464190388540141 1 0
-0.3465331894868054 1 0
-0.2466473401195857 1 0
-1 5 0 3
-25
-26
-27
-0.1100711180643749 0.9166596607351063 0
-0.07338074537641293 0.8333193214703032 0
-0.03669037268821046 0.7499789822049541 0
-1 6 0 8
-28
-29
-30
-31
-32
-33
-34
-35
-0.05149179445598822 0.5925676826130706 0
-0.1029835889118962 0.5184967222866699 0
-0.1544753833678243 0.4444257619602401 0
-0.2059671778237651 0.3703548016337923 0
-0.2574589722798281 0.2962838413071686 0
-0.3089507667360067 0.2222128809803785 0
-0.3604425611921816 0.1481419206535939 0
-0.4119343556483683 0.07407096032679217 0
-2 1 0 84
-36
-37
-38
-39
-40
-41
-42
-43
-44
-45
-46
-47
-48
-49
-50
-51
-52
-53
-54
-55
-56
-57
-58
-59
-60
-61
-62
-63
-64
-65
-66
-67
-68
-69
-70
-71
-72
-73
-74
-75
-76
-77
-78
-79
-80
-81
-82
-83
-84
-85
-86
-87
-88
-89
-90
-91
-92
-93
-94
-95
-96
-97
-98
-99
-100
-101
-102
-103
-104
-105
-106
-107
-108
-109
-110
-111
-112
-113
-114
-115
-116
-117
-118
-119
-0.2958604083772859 0.3779125235558248 0
-0.9170489619130662 0.6342263144098151 0
-0.4952131158506095 0.9135760699403515 0
-0.7486798353378267 0.1899126356545098 0
-0.123066269195724 0.7516284377586708 0
-0.398843997289799 0.2297706029024803 0
-0.1914841452302111 0.5250863010236176 0
-0.9177810706296824 0.4369587076892717 0
-0.8312448046640464 0.8474189616111325 0
-0.2983081248282825 0.9127303409426388 0
-0.5580920664369845 0.1182998099870039 0
-0.6887000767926235 0.9099416730617055 0
-0.3961268329415643 0.9140731297312352 0
-0.4452354243421484 0.8289568679789109 0
-0.5429139668610208 0.8274166831503765 0
-0.4936233968514193 0.7445239226727628 0
-0.5896364772446604 0.7426865367268549 0
-0.5408823976792568 0.6608925726988444 0
-0.6358632798126335 0.6593019073520822 0
-0.5869474609315031 0.5780940575804299 0
-0.3968859158423016 0.7456358057704324 0
-0.4922563164940828 0.5797605358949527 0
-0.5387681554514119 0.497080348400596 0
-0.6325690689158422 0.495904821987176 0
-0.5843916891691673 0.4155592542201335 0
-0.6779873544624372 0.4139662013578365 0
-0.726309298278548 0.494277578321845 0
-0.4912292252422175 0.4165523288108102 0
-0.5365810482928114 0.3355000825490801 0
-0.6852816676998132 0.7406467872116574 0
-0.730569909059229 0.6574793729447086 0
-0.7689548284328682 0.4105757920054097 0
-0.722929275213537 0.3350295701541776 0
-0.8269370850699298 0.3187639897622737 0
-0.7805798143213398 0.743343235029987 0
-0.8298954054390456 0.4877883917872448 0
-0.4436465794595489 0.3315637260249664 0
-0.3836028845107313 0.410608735622806 0
-0.8220807257969505 0.6530292564778822 0
-0.8692657716499751 0.7338908802730533 0
-0.5931282623436586 0.9124910954199575 0
-0.4914714363404978 0.2540855754977935 0
-0.5877252703314915 0.2552974817628334 0
-0.4642129036653649 0.1582963629378746 0
-0.3481303522893341 0.8287934341213055 0
-0.3022498599352635 0.7446152978129883 0
-0.3497197681707215 0.6636421939860775 0
-0.2589000618713882 0.6672458489825979 0
-0.3058119035980789 0.5893255017426672 0
-0.25151166623284 0.8263513208814195 0
-0.3952443288879547 0.5831843869047131 0
-0.3512774829448179 0.5076369703802126 0
-0.08839141565969492 0.6733512571582343 0
-0.7328251623565074 0.827895801810883 0
-0.4430063388437068 0.4985074278184796 0
-0.6398311725190902 0.1681143736057568 0
-0.834345999576188 0.2225497480225142 0
-0.3475366048037175 0.305807017337666 0
-0.9198396725461351 0.5303324882542779 0
-0.9193106474262465 0.3475804020648363 0
-0.1390142383104863 0.5984942222561735 0
-0.2425992422168039 0.4507534725024335 0
-0.6810658271232706 0.5766417987150884 0
-0.4452150720957505 0.6627796154465432 0
-0.6297826839798153 0.3343452771715426 0
-0.7740389604986767 0.5735597624649564 0
-0.2016255404130775 0.9155829532620882 0
-0.6394467887498616 0.824012146367313 0
-0.6830444844310272 0.2585155162313388 0
-0.2109812360889767 0.7474313627757466 0
-0.9249113999210528 0.8045032382443527 0
-0.9151941261815858 0.2574630825826384 0
-0.1602529739673592 0.8335224297280251 0
-0.4901041201884468 0.07675083066981332 0
-0.1728945660875086 0.6726530528222684 0
-0.8534289993071384 0.5740019041633093 0
-0.2206991050225435 0.5962689249355851 0
-0.7649391454339022 0.26593699684332 0
-0.8494376580426346 0.4007237314606419 0
-0.5447674825169194 0.1919391770036209 0
-0.2690976312251699 0.5195880146905167 0
-0.7695896036785943 0.9039057494253757 0
-0.9441074678470953 0.7122496975281268 0
-0.3127077125558917 0.4445574533926085 0
-$EndNodes
-$Elements
-7 236 1 236
-1 1 1 6
-1 1 7 
-2 7 8 
-3 8 9 
-4 9 10 
-5 10 11 
-6 11 2 
-1 2 1 7
-7 2 12 
-8 12 13 
-9 13 14 
-10 14 15 
-11 15 16 
-12 16 17 
-13 17 3 
-1 3 1 3
-14 3 18 
-15 18 19 
-16 19 4 
-1 4 1 6
-17 4 20 
-18 20 21 
-19 21 22 
-20 22 23 
-21 23 24 
-22 24 5 
-1 5 1 4
-23 5 25 
-24 25 26 
-25 26 27 
-26 27 6 
-1 6 1 9
-27 6 28 
-28 28 29 
-29 29 30 
-30 30 31 
-31 31 32 
-32 32 33 
-33 33 34 
-34 34 35 
-35 35 1 
-2 1 2 201
-36 73 36 93 
-37 44 70 75 
-38 91 78 115 
-39 69 92 107 
-40 95 69 107 
-41 72 73 93 
-42 46 91 115 
-43 74 37 75 
-44 75 37 118 
-45 41 72 93 
-46 70 44 89 
-47 44 75 106 
-48 78 91 104 
-49 34 35 79 
-50 71 43 94 
-51 72 41 77 
-52 24 5 102 
-53 105 85 108 
-54 43 71 114 
-55 5 25 102 
-56 41 34 79 
-57 40 105 108 
-58 36 73 119 
-59 47 76 103 
-60 57 86 90 
-61 86 87 90 
-62 87 73 90 
-63 46 79 109 
-64 89 47 103 
-65 79 46 115 
-66 47 20 76 
-67 79 35 109 
-68 56 82 99 
-69 86 57 99 
-70 45 24 102 
-71 56 49 80 
-72 56 80 81 
-73 20 21 76 
-74 38 48 49 
-75 49 48 80 
-76 82 86 99 
-77 38 22 48 
-78 56 81 82 
-79 40 27 88 
-80 48 45 80 
-81 77 41 79 
-82 80 45 85 
-83 51 49 56 
-84 82 81 83 
-85 50 49 51 
-86 23 24 45 
-87 51 56 99 
-88 23 45 48 
-89 58 57 90 
-90 46 8 91 
-91 38 49 50 
-92 52 51 53 
-93 4 20 47 
-94 82 83 84 
-95 50 51 52 
-96 81 80 85 
-97 21 38 76 
-98 54 53 55 
-99 27 6 88 
-100 52 53 54 
-101 55 57 58 
-102 55 53 57 
-103 64 63 72 
-104 72 63 73 
-105 74 66 101 
-106 7 8 46 
-107 57 53 99 
-108 21 22 38 
-109 18 19 44 
-110 9 39 91 
-111 86 84 87 
-112 60 58 63 
-113 22 23 48 
-114 8 9 91 
-115 64 72 77 
-116 65 54 66 
-117 60 63 64 
-118 39 10 92 
-119 55 58 59 
-120 70 66 74 
-121 91 39 104 
-122 70 74 75 
-123 63 58 90 
-124 66 54 98 
-125 10 11 92 
-126 9 10 39 
-127 59 58 60 
-128 82 84 86 
-129 59 62 98 
-130 55 59 98 
-131 52 54 65 
-132 62 67 71 
-133 38 50 76 
-134 67 68 69 
-135 85 45 102 
-136 73 63 90 
-137 76 50 103 
-138 54 55 98 
-139 65 66 70 
-140 12 13 95 
-141 31 32 36 
-142 62 61 67 
-143 59 61 62 
-144 53 51 99 
-145 30 31 97 
-146 59 60 61 
-147 67 61 68 
-148 13 14 43 
-149 29 30 42 
-150 32 33 93 
-151 26 27 40 
-152 14 15 94 
-153 28 29 96 
-154 15 16 37 
-155 6 28 88 
-156 31 36 97 
-157 68 61 100 
-158 33 34 41 
-159 13 43 95 
-160 36 32 93 
-161 42 30 97 
-162 64 77 78 
-163 43 14 94 
-164 29 42 96 
-165 66 98 101 
-166 62 71 101 
-167 60 64 100 
-168 15 37 94 
-169 88 28 96 
-170 33 41 93 
-171 61 60 100 
-172 98 62 101 
-173 50 52 103 
-174 65 70 89 
-175 64 78 100 
-176 52 65 103 
-177 17 106 118 
-178 81 85 105 
-179 68 100 104 
-180 17 3 106 
-181 100 78 104 
-182 69 95 114 
-183 65 89 103 
-184 18 44 106 
-185 40 88 110 
-186 71 94 111 
-187 1 7 109 
-188 11 2 107 
-189 35 1 109 
-190 2 12 107 
-191 85 102 108 
-192 92 69 113 
-193 102 25 108 
-194 25 26 108 
-195 83 81 105 
-196 106 75 118 
-197 3 18 106 
-198 73 87 119 
-199 105 40 110 
-200 4 47 117 
-201 71 67 114 
-202 37 74 111 
-203 104 39 113 
-204 92 11 107 
-205 19 4 117 
-206 68 104 113 
-207 12 95 107 
-208 37 16 118 
-209 95 43 114 
-210 74 101 111 
-211 26 40 108 
-212 7 46 109 
-213 84 83 112 
-214 89 44 117 
-215 101 71 111 
-216 83 110 112 
-217 87 84 116 
-218 83 105 110 
-219 47 89 117 
-220 84 112 116 
-221 88 96 110 
-222 87 116 119 
-223 94 37 111 
-224 39 92 113 
-225 110 96 112 
-226 69 68 113 
-227 77 79 115 
-228 67 69 114 
-229 96 42 112 
-230 78 77 115 
-231 112 42 116 
-232 44 19 117 
-233 42 97 116 
-234 16 17 118 
-235 116 97 119 
-236 97 36 119 
-$EndElements
diff --git a/test_ring_demo_perf.py b/test_ring_demo_perf.py
index 49ee937..4095d7d 100644
--- a/test_ring_demo_perf.py
+++ b/test_ring_demo_perf.py
@@ -1,11 +1,13 @@
 import os
+import pickle
 import time
 from types import SimpleNamespace
+
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
 import yaml
-import pickle
+
 import UM2N
 
 print("Setting up solver.")
diff --git a/tests/dataset_integrity_check.ipynb b/tests/dataset_integrity_check.ipynb
index 6983e18..19af45d 100644
--- a/tests/dataset_integrity_check.ipynb
+++ b/tests/dataset_integrity_check.ipynb
@@ -9,12 +9,14 @@
     "%load_ext autoreload\n",
     "%autoreload 2\n",
     "\n",
-    "import UM2N\n",
+    "import warnings\n",
+    "\n",
     "# import glob\n",
     "import torch\n",
     "from torch_geometric.utils import index_to_mask\n",
     "\n",
-    "import warnings\n",
+    "import UM2N\n",
+    "\n",
     "warnings.filterwarnings('ignore')\n",
     "device = torch.device('cuda' if torch.cuda.is_available()\n",
     "else 'cpu')"
diff --git a/tests/test_unstructured_mesh.py b/tests/test_unstructured_mesh.py
new file mode 100644
index 0000000..aed93b0
--- /dev/null
+++ b/tests/test_unstructured_mesh.py
@@ -0,0 +1,118 @@
+"""
+Unit tests for the generate_mesh mesh generator module.
+"""
+
+import os
+
+import numpy as np
+import pytest
+import ufl
+from firedrake.assemble import assemble
+from firedrake.bcs import DirichletBC
+from firedrake.constant import Constant
+
+from UM2N.generator.unstructured_mesh import (
+    UnstructuredRandomPolygonalMeshGenerator,
+    UnstructuredSquareMeshGenerator,
+)
+
+
+@pytest.fixture(params=[1, 2, 3, 4])
+def num_elem_bnd(request):
+    return request.param
+
+
+@pytest.fixture(params=[1, 10, 0.2, np.pi])
+def scale(request):
+    return request.param
+
+
+@pytest.fixture(params=[1, 2], ids=["delaunay", "frontal"])
+def mesh_algorithm(request):
+    return request.param
+
+
+@pytest.fixture(
+    params=[
+        UnstructuredRandomPolygonalMeshGenerator,
+        UnstructuredSquareMeshGenerator,
+    ]
+)
+def generator(request):
+    return request.param
+
+
+def generate_mesh(generator, mesh_algorithm, scale=1.0, **kwargs):
+    """
+    Utility mesh generator function for testing purposes.
+    """
+    mesh_gen = generator(mesh_type=mesh_algorithm, scale=scale)
+    kwargs.setdefault("remove_file", True)
+    mesh = mesh_gen.generate_mesh(**kwargs)
+    mesh.init()
+    return mesh
+
+
+def test_file_removal():
+    """
+    Test that the remove_file keyword argument works as expected.
+    """
+    output_filename = "./tmp.msh"
+    assert not os.path.exists(output_filename)
+    generate_mesh(
+        UnstructuredSquareMeshGenerator,
+        1,
+        res=1.0,
+        output_filename=output_filename,
+        remove_file=False,
+    )
+    assert os.path.exists(output_filename)
+    os.remove(output_filename)
+    assert not os.path.exists(output_filename)
+    generate_mesh(
+        UnstructuredSquareMeshGenerator, 1, res=1.0, output_filename=output_filename
+    )
+    assert not os.path.exists(output_filename)
+
+
+def test_boundary_segments(generator):
+    """
+    Check that the boundary segments are tagged with integers counting from 1.
+    """
+    mesh = generate_mesh(generator, 1, res=1.0)
+    boundary_ids = mesh.exterior_facets.unique_markers
+    assert set(boundary_ids) == set(range(1, len(boundary_ids) + 1))
+
+
+def test_num_points_boundaries_square(num_elem_bnd, mesh_algorithm):
+    """
+    Check that the numbers of points on each boundary segment of a unit square mesh are
+    as expected.
+    """
+    mesh = generate_mesh(UnstructuredSquareMeshGenerator, 1, res=1.0 / num_elem_bnd)
+    boundary_ids = mesh.exterior_facets.unique_markers
+    for boundary_id in boundary_ids:
+        dbc = DirichletBC(mesh.coordinates.function_space(), 0, boundary_id)
+        assert len(dbc.nodes) == num_elem_bnd + 1
+
+
+def test_area_squaremesh(num_elem_bnd, mesh_algorithm, scale):
+    """
+    Check that the area of a square mesh is equal to the scale factor squared.
+    """
+    mesh = generate_mesh(
+        UnstructuredSquareMeshGenerator, 1, res=1.0 / num_elem_bnd, scale=scale
+    )
+    assert np.isclose(assemble(Constant(1.0, domain=mesh) * ufl.dx), scale**2)
+
+
+def test_num_cells_with_res_and_scale(generator, num_elem_bnd, mesh_algorithm):
+    """
+    Check that doubling or halving the overall resolution doesn't affect the number of
+    cells for the square mesh, so long as the resolution is changed accordingly.
+    """
+    generator = UnstructuredSquareMeshGenerator
+    mesh1 = generate_mesh(generator, mesh_algorithm, res=1.0 / num_elem_bnd)
+    mesh2 = generate_mesh(generator, mesh_algorithm, res=2.0 / num_elem_bnd, scale=2.0)
+    meshp5 = generate_mesh(generator, mesh_algorithm, res=0.5 / num_elem_bnd, scale=0.5)
+    assert np.allclose((mesh2.num_cells(), meshp5.num_cells()), mesh1.num_cells())