src/scripts/model_maker.py

@@ -24,6 +24,7 @@
 #  The script requires python3.
 
 import math
+import numpy as np
 import os
 import pdb
 import random
@@ -36,7 +37,7 @@ except ModuleNotFoundError as ex:
     print("WARNING: pyvista not found!")
     allow_3d = False
 
-from pathlib import PurePath
+from pathlib import Path
 from sys import argv
 
 ## \brief Main execution code
@@ -69,7 +70,7 @@ def interpolate_constants(sconf):
     for i in range(sconf['configurations']):
         for j in range(sconf['nshl'][i]):
             file_idx = sconf['dielec_id'][i][j]
-            dielec_path = PurePath(sconf['dielec_path'], sconf['dielec_file'][int(file_idx) - 1])
+            dielec_path = Path(sconf['dielec_path'], sconf['dielec_file'][int(file_idx) - 1])
             file_name = str(dielec_path)
             dielec_file = open(file_name, 'r')
             wavelengths = []
@@ -149,7 +150,7 @@ def load_model(model_file):
             make_3d = False
         # Create the sconf dict
         sconf = {
-            'out_file': PurePath(
+            'out_file': Path(
                 model['input_settings']['input_folder'],
                 model['input_settings']['spheres_file']
             )
@@ -316,7 +317,7 @@ def load_model(model_file):
             print("ERROR: %s is not a recognized polarization state."%str_polar)
             return (None, None)
         gconf = {
-            'out_file': PurePath(
+            'out_file': Path(
                 model['input_settings']['input_folder'],
                 model['input_settings']['geometry_file']
             )
@@ -365,9 +366,13 @@ def load_model(model_file):
             if (len_vec_x == 0):
                 # Generate random cluster
                 rnd_seed = int(model['system_settings']['rnd_seed'])
-                # random_aggregate() checks internally whether application is INCLUSION
                 max_rad = float(model['particle_settings']['max_rad'])
-                random_aggregate(sconf, gconf, rnd_seed, max_rad)
+                # random_aggregate() checks internally whether application is INCLUSION
+                #random_aggregate(sconf, gconf, rnd_seed, max_rad)
+                check = random_compact(sconf, gconf, rnd_seed, max_rad)
+                if (check != 0):
+                    print("INFO: stopping with exit code %d."%check)
+                    exit(check)
             else:
                 if (len(model['geometry_settings']['x_coords']) != gconf['nsph']):
                     print("ERROR: coordinate vectors do not match the number of spheres!")
@@ -404,7 +409,7 @@ def match_grid(sconf):
             layers += 1
         for j in range(layers):
             file_idx = sconf['dielec_id'][i][j]
-            dielec_path = PurePath(sconf['dielec_path'], sconf['dielec_file'][int(file_idx) - 1])
+            dielec_path = Path(sconf['dielec_path'], sconf['dielec_file'][int(file_idx) - 1])
             file_name = str(dielec_path)
             dielec_file = open(file_name, 'r')
             wavelengths = []
@@ -522,7 +527,10 @@ def print_help():
 #  \param seed: `int` Seed for the random sequence generation
 #  \param max_rad: `float` Maximum allowed radial extension of the aggregate
 #  \param max_attempts: `int` Maximum number of attempts to place a particle in any direction
+#  \return result: `int` Function exit code (0 for success, otherwise number of
+#  spheres that could not be placed)
 def random_aggregate(scatterer, geometry, seed, max_rad, max_attempts=100):
+    result = 0
     random.seed(seed)
     nsph = scatterer['nsph']
     vec_thetas = [0.0 for i in range(nsph)]
@@ -545,6 +553,7 @@ def random_aggregate(scatterer, geometry, seed, max_rad, max_attempts=100):
         is_placed = False
         while (not is_placed):
             if (attempts > max_attempts):
+                result += 1
                 print("WARNING: could not place sphere %d in allowed radius!"%i)
                 break # while(not is_placed)
             vec_thetas[i] = math.pi * random.random()
@@ -618,7 +627,122 @@ def random_aggregate(scatterer, geometry, seed, max_rad, max_attempts=100):
         geometry['vec_sph_y'][sph_index] = sphere['y']
         geometry['vec_sph_z'][sph_index] = sphere['z']
         sph_index += 1
-# end random_aggregate()
+    return result
+
+## \brief Generate a random compact cluster from YAML configuration options.
+#
+#  This function generates a random aggregate of spheres using the maximum
+#  compactness packaging to fill a spherical volume with given maximum radius,
+#  then it proceeds by subtracting random spheres from the outer layers, until
+#  the aggregate is reduced to the desired number of spheres. The function
+#  can only be used if all sphere types have the same radius. The result of the
+#  generated model is directly saved in the parameters of the scatterer and
+#  geometry configuration dictionaries.
+#
+#  \param scatterer: `dict` Scatterer configuration dictionary (gets modified)
+#  \param geometry: `dict` Geometry configuration dictionary (gets modified)
+#  \param seed: `int` Seed for the random sequence generation
+#  \param max_rad: `float` Maximum allowed radial extension of the aggregate
+#  \return result: `int` Function exit code (0 for success, otherwise error code)
+def random_compact(scatterer, geometry, seed, max_rad):
+    result = 0
+    random.seed(seed)
+    nsph = scatterer['nsph']
+    n_types = scatterer['configurations']
+    if (0 in scatterer['vec_types']):
+        tincrement = 1 if scatterer['application'] != "INCLUSION" else 2
+        for ti in range(nsph):
+            itype = tincrement + int(n_types * random.random())
+            scatterer['vec_types'][ti] = itype
+    if (max(scatterer['ros']) != min(scatterer['ros'])):
+        result = 1
+    else:
+        radius = scatterer['ros'][0]
+        x_centers = np.arange(-1.0 * max_rad + radius, max_rad, 2.0 * radius)
+        y_centers = np.arange(-1.0 * max_rad + radius, max_rad, math.sqrt(3.0) * radius)
+        z_centers = np.arange(-1.0 * max_rad + radius, max_rad, math.sqrt(3.0) * radius)
+        x_offset = radius
+        y_offset = radius / math.sqrt(3.0)
+        z_offset = 0.0
+        tmp_spheres = []
+        n_cells = len(x_centers) * len(y_centers) * len(z_centers)
+        print("INFO: the cubic space would contain %d spheres."%n_cells)
+        n_max_spheres = int(max_rad * max_rad * max_rad / (radius * radius * radius) * 0.74)
+        print("INFO: the maximum radius allows for %d spheres."%n_max_spheres)
+        for zi in range(len(z_centers)):
+            if (y_offset == 0.0):
+                y_offset = radius / math.sqrt(3.0)
+            else:
+                y_offset = 0.0
+            for yi in range(len(y_centers)):
+                if (x_offset == 0.0):
+                    x_offset = radius
+                else:
+                    x_offset = 0.0
+                for xi in range(len(x_centers)):
+                    x = x_centers[xi] + x_offset
+                    y = y_centers[yi] + y_offset
+                    z = z_centers[zi]
+                    extent = radius + math.sqrt(x * x + y * y + z * z)
+                    if (extent < max_rad):
+                        tmp_spheres.append({
+                            'itype': 1,
+                            'x': x,
+                            'y': y,
+                            'z': z
+                        })
+        #tmp_spheres = [{'itype': 1, 'x': 0.0, 'y': 0.0, 'z': 0.0}]
+        current_n = len(tmp_spheres)
+        print("INFO: before erosion there are %d spheres in use."%current_n)
+        rho = 2.0 * max_rad
+        while (current_n > nsph):
+            theta = 2.0 * math.pi * random.random()
+            phi = math.pi * random.random()
+            x0 = rho * math.sin(theta) * math.cos(phi)
+            y0 = rho * math.sin(theta) * math.sin(phi)
+            z0 = rho * math.cos(theta)
+            closest_index = 0
+            minimum_distance = 1000.0 * max_rad
+            for di in range(len(tmp_spheres)):
+                x1 = tmp_spheres[di]['x']
+                if (x1 == max_rad):
+                    continue
+                y1 = tmp_spheres[di]['y']
+                z1 = tmp_spheres[di]['z']
+                distance = math.sqrt(
+                    (x1 - x0) * (x1 - x0)
+                    + (y1 - y0) * (y1 - y0)
+                    + (z1 - z0) * (z1 - z0)
+                )
+                if (distance < minimum_distance):
+                    closest_index = di
+                    minimum_distance = distance
+            tmp_spheres[closest_index]['x'] = max_rad
+            current_n -= 1
+        vec_spheres = []
+        sph_index = 0
+        for ti in range(len(tmp_spheres)):
+            sphere = tmp_spheres[ti]
+            if (sphere['x'] < max_rad):
+                sphere['itype'] = scatterer['vec_types'][sph_index]
+                sph_index += 1
+                vec_spheres.append(sphere)
+        #pl = pv.Plotter()
+        #for si in range(len(vec_spheres)):
+        #    x = vec_spheres[si]['x'] / max_rad
+        #    y = vec_spheres[si]['y'] / max_rad
+        #    z = vec_spheres[si]['z'] / max_rad
+        #    mesh = pv.Sphere(radius / max_rad, (x, y, z))
+        #    pl.add_mesh(mesh)
+        #pl.export_obj("scene.obj")
+    sph_index = 0
+    for sphere in sorted(vec_spheres, key=lambda item: item['itype']):
+        scatterer['vec_types'][sph_index] = sphere['itype']
+        geometry['vec_sph_x'][sph_index] = sphere['x']
+        geometry['vec_sph_y'][sph_index] = sphere['y']
+        geometry['vec_sph_z'][sph_index] = sphere['z']
+        sph_index += 1
+    return result
     
 ## \brief Write the geometry configuration dictionary to legacy format.
 #
@@ -784,6 +908,9 @@ def write_legacy_sconf(conf):
 #  \param geometry: `dict` Geometry configuration dictionary (gets modified)
 #  \param max_rad: `float` Maximum allowed radial extension of the aggregate
 def write_obj(scatterer, geometry, max_rad):
+    out_dir = scatterer['out_file'].absolute().parent
+    out_model_path = Path(out_dir, "model.obj")
+    out_material_path = Path(out_dir, "model.mtl")
     color_strings = [
         "1.0 1.0 1.0\n", # white
         "1.0 0.0 0.0\n", # red
@@ -793,9 +920,9 @@ def write_obj(scatterer, geometry, max_rad):
     color_names = [
         "white", "red", "blue", "green"
     ]
-    mtl_file = open("model.mtl", "w")
+    mtl_file = open(str(out_material_path), "w")
     for mi in range(len(color_strings)):
-        mtl_line = "newmtl mtl{0:d}\n".format(mi)
+        mtl_line = "newmtl "  + color_names[mi] + "\n"
         mtl_file.write(mtl_line)
         color_line = color_strings[mi]
         mtl_file.write("   Ka " + color_line)
@@ -808,28 +935,43 @@ def write_obj(scatterer, geometry, max_rad):
     pl = pv.Plotter()
     for si in range(scatterer['nsph']):
         sph_type_index = scatterer['vec_types'][si]
-        color_by_name = color_names[sph_type_index]
+        # color_index = 1 + (sph_type_index % (len(color_strings) - 1))
+        # color_by_name = color_names[sph_type_index]
         radius = scatterer['ros'][sph_type_index - 1] / max_rad
         x = geometry['vec_sph_x'][si] / max_rad
         y = geometry['vec_sph_y'][si] / max_rad
         z = geometry['vec_sph_z'][si] / max_rad
         mesh = pv.Sphere(radius, (x, y, z))
-        mesh.save("tmp_mesh.obj")
-        pl.add_mesh(mesh) #, color=color_by_name)
-        mesh_name = "sphere_{0:04d}.obj".format(si)
-        in_obj_file = open("tmp_mesh.obj", "r")
-        out_obj_file = open(mesh_name, "w")
-        in_line = in_obj_file.readline()
-        out_obj_file.write(in_line)
-        out_obj_file.write("mtllib model.mtl\n")
-        out_obj_file.write("usemtl mtl{0:d}\n".format(sph_type_index))
-        while (in_line != ""):
-            in_line = in_obj_file.readline()
-            out_obj_file.write(in_line)
-        in_obj_file.close()
-        out_obj_file.close()
-    pl.export_obj("model.obj")
-    os.remove("tmp_mesh.obj")
+        pl.add_mesh(mesh, color=None)
+    pl.export_obj(str(Path(str(out_dir), "TMP_MODEL.obj")))
+    tmp_model_file = open(str(Path(str(out_dir), "TMP_MODEL.obj")), "r")
+    out_model_file = open(str(Path(str(out_dir), "model.obj")), "w")
+    sph_index = 0
+    sph_type_index = 0
+    old_sph_type_index = 0
+    str_line = tmp_model_file.readline()
+    while (str_line != ""):
+        if (str_line.startswith("mtllib")):
+            str_line = "mtllib model.mtl\n"
+        elif (str_line.startswith("g ")):
+            sph_index += 1
+            sph_type_index = scatterer['vec_types'][sph_index - 1]
+            if (sph_type_index == old_sph_type_index):
+                str_line = tmp_model_file.readline()
+                str_line = tmp_model_file.readline()
+            else:
+                old_sph_type_index = sph_type_index
+                color_index = sph_type_index % (len(color_names) - 1)
+                str_line = "g grp{0:04d}\n".format(sph_type_index)
+                out_model_file.write(str_line)
+                str_line = tmp_model_file.readline()
+                str_line = "usemtl {0:s}\n".format(color_names[color_index])
+        out_model_file.write(str_line)
+        str_line = tmp_model_file.readline()
+    out_model_file.close()
+    tmp_model_file.close()
+    os.remove(str(Path(str(out_dir), "TMP_MODEL.obj")))
+    os.remove(str(Path(str(out_dir), "TMP_MODEL.mtl")))
 
 ## \brief Exit code (0 for success)
 exit_code = main()