SphericalEllipseOverlap/packing_simulation.py

import os

os.environ["OPENBLAS_NUM_THREADS"] = "1"  # mitigates the clash between numpy and multiprocessing  parallelization

import numpy as np
import ctypes
import multiprocessing
from functools import partial
import time
from scipy.optimize import minimize as minim, OptimizeResult
from dataclasses import dataclass, field
from abc import ABC, abstractmethod


def ephi(phi):
    u = np.zeros((len(phi), 3))
    u[:, 0] = -np.sin(phi)
    u[:, 1] = np.cos(phi)
    return u


def etheta(phi, theta):
    u = np.zeros((len(phi), 3))
    u[:, 0] = np.cos(phi) * np.cos(theta)
    u[:, 1] = np.sin(phi) * np.cos(theta)
    u[:, 2] = -np.sin(theta)
    return u


def sph_to_cart(sph):
    """Transformation form spherical coordinates phi-theta on a sphere to cartesian coordinates."""
    cart = np.zeros((len(sph), 3))
    cart[:, 0] = np.cos(sph[:, 0]) * np.sin(sph[:, 1])
    cart[:, 1] = np.sin(sph[:, 0]) * np.sin(sph[:, 1])
    cart[:, 2] = np.cos(sph[:, 1])
    return cart


def vector_orientations(sph, xi):
    """
    Calculates vector_orientations of ellipse orientations given their positions in spherical coordinates
    and orientation angles.
    """
    cosxi = np.cos(xi)[:, np.newaxis]
    sinxi = np.sin(xi)[:, np.newaxis]
    basis_theta = etheta(sph[:, 0], sph[:, 1])
    basis_phi = ephi(sph[:, 0])
    return cosxi * basis_theta + sinxi * basis_phi


clib = ctypes.cdll.LoadLibrary('./overlap_algorithm.so')
clib.contact_mono_pack_multi.argtypes = [
    ctypes.c_int,
    ctypes.c_double,
    ctypes.c_double,
    np.ctypeslib.ndpointer(dtype=np.float64, ndim=2, flags='C_CONTIGUOUS'),
    np.ctypeslib.ndpointer(dtype=np.float64, ndim=2, flags='C_CONTIGUOUS'),
    ctypes.c_double
    ]
clib.contact_mono_pack_multi.restype = ctypes.c_double

clib.contact_bidisp_pack_multi.argtypes = [
    ctypes.c_int,
    np.ctypeslib.ndpointer(dtype=np.float64, ndim=1, flags='C_CONTIGUOUS'),
    np.ctypeslib.ndpointer(dtype=np.float64, ndim=1, flags='C_CONTIGUOUS'),
    np.ctypeslib.ndpointer(dtype=np.float64, ndim=2, flags='C_CONTIGUOUS'),
    np.ctypeslib.ndpointer(dtype=np.float64, ndim=2, flags='C_CONTIGUOUS'),
    ctypes.c_double,
    ctypes.c_double
    ]
clib.contact_bidisp_pack_multi.restype = ctypes.c_double


def monodisp_config_energy(inv2a, inv2b, coords, orients, dist_lim) -> float:
    if inv2a >= 1:
        n = len(coords)
        return clib.contact_mono_pack_multi(n, inv2a, inv2b, coords, orients, dist_lim)
    raise ValueError("Semi-major axis should be smaller that pi/2, otherwise the ellipse becomes inverted.")


def bidisp_config_energy(inv2a_array, inv2b_array, coords, orients, dist_lim, omega) -> float:
    if inv2a_array[0] >= 1:
        n = len(coords)
        return clib.contact_bidisp_pack_multi(n, inv2a_array, inv2b_array, coords, orients, dist_lim, omega)
    raise ValueError("Semi-major axis should be smaller that pi/2, otherwise the ellipse becomes inverted.")


class EllipseMinim(ABC):

    """Base class fir optimization of ellipse configurations on a sphere."""

    @abstractmethod
    def change_params(self, new_alpha: float) -> None:
        pass

    @abstractmethod
    def evaluate_energy(self) -> float:
        pass

    @abstractmethod
    def export_config(self, folder: str, idx: int) -> None:
        pass

    @abstractmethod
    def minim(self, eps, gtol) -> OptimizeResult:
        """
        Method for finding optimal configuration with BFGS energy minimization method.
        :param eps: Absolute step size used for numerical approximation of the jacobian in the BFGS method.
        :param gtol: BFGS minimization is terminated after all gradient components are smaller than gtol (inf norm)
        """
        pass


class EllipseMinimMonodisp(EllipseMinim):

    def __init__(self, n, alpha, epsilon, initial: np.ndarray = None):
        """
        :param n: number of spherical ellipses
        :param alpha: ellipse size (geodesic semi-major axis)
        :param epsilon: ellipse aspect ratio
        """
        self.n = n
        self.alpha = alpha
        self.beta = alpha / epsilon
        self.epsilon = epsilon

        self.config_xi = initial
        if self.config_xi is None:
            np.random.seed()
            self.config_xi = np.array([2 * np.pi * np.random.random(self.n),
                                       np.arccos(2 * np.random.random(self.n) - 1),
                                       2 * np.pi * np.random.random(self.n)]).flatten()

        self.inv2a = 1 / np.sin(self.alpha) ** 2
        self.inv2b = 1 / np.sin(self.beta) ** 2
        self.dist_lim = np.cos(2 * self.alpha)

    def change_params(self, new_alpha):
        self.alpha = new_alpha
        self.beta = new_alpha / self.epsilon
        self.inv2a = 1 / np.sin(self.alpha) ** 2
        self.inv2b = 1 / np.sin(self.beta) ** 2
        self.dist_lim = np.cos(2 * self.alpha)

    def evaluate_energy(self):
        sph = self.config_xi[:2 * self.n].reshape((2, self.n)).T
        positions = sph_to_cart(sph)
        orientations = vector_orientations(sph, self.config_xi[2 * self.n:])
        return monodisp_config_energy(self.inv2a, self.inv2b, positions, orientations, self.dist_lim)

    def export_config(self, folder: str, idx: int):
        sph = self.config_xi[:2 * self.n].reshape((2, self.n)).T
        positions = sph_to_cart(sph)
        orientations = vector_orientations(sph, self.config_xi[2 * self.n:])
        np.savetxt(folder + f'/coord{idx}.dat', positions)
        np.savetxt(folder + f'/orient{idx}.dat', orientations)

    def minim(self, eps=1.4901161193847656e-08, gtol=1e-5):

        def energy(xi):
            sph = xi[:2 * self.n].reshape((2, self.n)).T
            positions = sph_to_cart(sph)
            orientations = vector_orientations(sph, xi[2 * self.n:])
            return monodisp_config_energy(self.inv2a, self.inv2b, positions, orientations, self.dist_lim)

        result = minim(energy, self.config_xi, method='BFGS', options={'gtol': gtol, 'eps': eps})
        self.config_xi = result.x

        return result


class EllipseMinimBidisp(EllipseMinim):

    def __init__(self, n, alpha1, alpha2, epsilon, num_frac=0.5, initial=None):
        """
        :param n: number of spherical ellipses
        :param alpha1: geodesic semi-major axis for the 1st size of ellipses
        :param alpha2: geodesic semi-major axis for the 2nd size of ellipses
        :param epsilon: ellipse aspect ratio
        :param num_frac: number fraction of larger ellipses
        """

        alpha1, alpha2 = max(alpha1, alpha2), min(alpha1, alpha2)
        if epsilon < 1:
            epsilon = 1 / epsilon

        self.n = n
        self.alpha1 = alpha1
        self.alpha2 = alpha2
        self.beta1 = alpha1 / epsilon
        self.beta2 = alpha2 / epsilon
        self.num_frac = num_frac
        self.size_ratio = alpha1 / alpha2
        self.epsilon = epsilon

        self.config_xi = initial
        if self.config_xi is None:
            np.random.seed()
            self.config_xi = np.array([2 * np.pi * np.random.random(self.n),
                                       np.arccos(2 * np.random.random(self.n) - 1),
                                       2 * np.pi * np.random.random(self.n)]).flatten()

        self.num_large = int(n * num_frac)

        self.inv2a = np.zeros(self.n)
        self.inv2b = np.zeros(self.n)
        self.inv2a[:self.num_large] = 1 / np.sin(self.alpha1) ** 2
        self.inv2a[self.num_large:] = 1 / np.sin(self.alpha2) ** 2
        self.inv2b[:self.num_large] = 1 / np.sin(self.beta1) ** 2
        self.inv2b[self.num_large:] = 1 / np.sin(self.beta2) ** 2

        self.dist_lim = np.cos(2 * self.alpha1)
        self.omega = 1 / np.sin(self.alpha1) ** 2

    def change_params(self, new_alpha1):
        self.alpha1 = new_alpha1
        self.alpha2 = new_alpha1 / self.size_ratio
        self.beta1 = self.alpha1 / self.epsilon
        self.beta2 = self.alpha2 / self.epsilon

        self.inv2a[:self.num_large] = 1 / np.sin(self.alpha1) ** 2
        self.inv2a[self.num_large:] = 1 / np.sin(self.alpha2) ** 2
        self.inv2b[:self.num_large] = 1 / np.sin(self.beta1) ** 2
        self.inv2b[self.num_large:] = 1 / np.sin(self.beta2) ** 2

        self.dist_lim = np.cos(2 * self.alpha1)

    def evaluate_energy(self):
        sph = self.config_xi[:2 * self.n].reshape((2, self.n)).T
        positions = sph_to_cart(sph)
        orientations = vector_orientations(sph, self.config_xi[2 * self.n:])
        return bidisp_config_energy(self.inv2a, self.inv2b, positions, orientations, self.dist_lim, self.omega)

    def export_config(self, folder: str, idx: int):
        sph = self.config_xi[:2 * self.n].reshape((2, self.n)).T
        positions = sph_to_cart(sph)
        orientations = vector_orientations(sph, self.config_xi[2 * self.n:])
        np.savetxt(folder + f'/coord{idx}.dat', positions)
        np.savetxt(folder + f'/orient{idx}.dat', orientations)

    def minim(self, eps=1.4901161193847656e-08, gtol=1e-5):

        def energy_eig(xi):
            sph = xi[:2 * self.n].reshape((2, self.n)).T
            positions = sph_to_cart(sph)
            orientations = vector_orientations(sph, xi[2 * self.n:])
            return bidisp_config_energy(self.inv2a, self.inv2b, positions, orientations, self.dist_lim, self.omega)

        result = minim(energy_eig, self.config_xi, method='BFGS',  options={'gtol': gtol, 'eps': eps})
        self.config_xi = result.x

        return result


def ellipse_minim_factory(n, alpha1, alpha2, epsilon, num_frac=0.5) -> EllipseMinim:
    if alpha1 == alpha2:
        return EllipseMinimMonodisp(n, alpha1, epsilon)
    return EllipseMinimBidisp(n, alpha1, alpha2, epsilon, num_frac=num_frac)


class PackingSimulation:

    def __init__(self, n: int, epsilon: float, data_folder: str = None, size_ratio: float = 1., num_frac: float = 0.5):
        self.n = n
        self.epsilon = epsilon
        self.num_frac = num_frac

        self.size_ratio = size_ratio
        if self.size_ratio < 1:
            self.size_ratio = 1 / self.size_ratio

        self.data_folder = data_folder
        if self.data_folder is None:
            self.data_folder = os.getcwd()

        if size_ratio == 1:
            self.folder = self.data_folder + f'/packing_results/monodisp/n{n}ratio{int(1000 * epsilon)}'
        else:
            self.folder = self.data_folder + \
            f'/packing_results/bidisp{int(100 * size_ratio)}frac{int(100 * num_frac)}/n{n}ratio{int(1000 * epsilon)}'

        try:
            existing = np.loadtxt(self.folder + '/ellipse_sizes.dat')
            self.existing_idx = len(existing)
        except OSError:
            self.existing_idx = 0
        except TypeError:
            self.existing_idx = 1

    def run(self, num_configs, tol=1e-8, n_jobs=8):
        os.makedirs(self.folder, exist_ok=True)

        if self.existing_idx != 0:
            print(f'{self.existing_idx} configurations for given parameters already exist.')

        threads = multiprocessing.Pool(n_jobs)
        ellipse_sizes = threads.map(partial(ellipse_packing, n=self.n, epsilon=self.epsilon, tol=tol,
                                            size_ratio=self.size_ratio, num_frac=self.num_frac,
                                            folder=self.folder, existing_idx=self.existing_idx),
                                    range(num_configs))
        threads.close()

        with open(self.folder + '/ellipse_sizes.dat', 'ab') as f:
            np.savetxt(f, ellipse_sizes)

        self.existing_idx += num_configs


@dataclass
class PackingSimData:

    minim_iter: list = field(init=False, default_factory=list)
    minim_nfev: list = field(init=False, default_factory=list)
    ellipse_scaling: list = field(init=False, default_factory=list)
    scaling_step: list = field(init=False, default_factory=list)
    energy: list = field(init=False, default_factory=list)

    def append_all(self, minim_result: OptimizeResult, ellipse_scaling: float, scaling_step: float):
        self.minim_iter.append(minim_result.nit)
        self.minim_nfev.append(minim_result.nfev)
        self.ellipse_scaling.append(ellipse_scaling)
        self.scaling_step.append(scaling_step)
        self.energy.append(minim_result.fun)

    def export(self, folder: str, idx: int):
        minim_iter = np.array(self.minim_iter)
        minim_nfev = np.array(self.minim_nfev)
        ellipse_scaling = np.array(self.ellipse_scaling)
        scaling_step = np.array(self.scaling_step)
        energy = np.array(self.energy)
        np.savetxt(folder + f'/iter_data{idx}.dat',
                   np.vstack((ellipse_scaling, scaling_step, minim_iter, minim_nfev, energy)).T)


def ellipse_packing(idx, epsilon: float, n: int, tol: float, size_ratio: float,
                    num_frac: float, folder: str, existing_idx: int = 0):

    scaling = np.sqrt(2 * epsilon / n)  # we start with scaling for packing fraction cca phi<0.5
    scaling_step0 = scaling / 100
    scaling_step = scaling_step0

    inflate_deflate = 1
    counter = 0
    n_iter = 0

    # default values
    eps = 1.4901161193847656e-08
    gtol = 1e-5

    packing_sim_data = PackingSimData()

    # random initialization of ellipse positions and orientations
    ellipse_minim = ellipse_minim_factory(n, scaling, scaling / size_ratio, epsilon, num_frac=num_frac)
    minim_result = ellipse_minim.minim(eps=eps, gtol=gtol)
    packing_sim_data.append_all(minim_result, scaling, scaling_step)

    if minim_result.fun > 1e-4:
        raise ValueError("Check initialization (packing fraction), the first minimization should get rid of overlaps.")

    # first: lowering of eps and gtol as long as minimization does not reach the desired accuracy
    while True:
        if minim_result.fun < 1e-12 or eps < 2 * 1e-10:
            break

        eps = eps / 2
        print(f'Lowering eps at first minim for idx {idx}, now {eps}')
        minim_result = ellipse_minim.minim(eps=eps, gtol=gtol)
        packing_sim_data.append_all(minim_result, scaling, scaling_step)

        # lowering the desired accuracy of gradient calculation if lowering eps does not help
        delta_en = packing_sim_data.energy[-2] - packing_sim_data.energy[-1]
        if delta_en < packing_sim_data.energy[-2] / 10 and gtol > 2 * 1e-7:
            gtol = gtol / np.sqrt(10)
            print(f'Lowering gtol at first minim for idx {idx}, now {gtol}')

    while True:
        if n_iter == 300:
            print('Exceeded max number of scaling iterations (n_iter_max=300).')
            break

        energy = ellipse_minim.evaluate_energy()
        if tol < energy < 2 * tol:
            break
        elif energy < tol and inflate_deflate == -1:
            scaling_step = scaling_step / 2
            inflate_deflate = 1  # ellipse size is increased
            counter = 0
        elif energy > 2 * tol and inflate_deflate == 1:
            scaling_step = scaling_step / 2
            inflate_deflate = -1  # ellipse size is decreased
            counter = 0
        # if there is no change in scaling step for 8 minimizations, we increase it
        elif counter == 8 and scaling_step < scaling_step0:
            scaling_step = 2 * scaling_step
            counter = 0

        # additional adjustment of eps and gtol:
        # we mandate strict energy accuracy during the first 10 iterations (there should still be no overlaps)
        if n_iter <= 10:
            if energy > 2 * 1e-12 and eps > 2 * 1e-10:
                eps = eps / 2
                print(f'Lowering eps at {n_iter + 1}. minim for idx {idx}, now {eps}')
                minim_result = ellipse_minim.minim(eps=eps, gtol=gtol)
                packing_sim_data.append_all(minim_result, scaling, scaling_step)

                delta_en = packing_sim_data.energy[-2] - packing_sim_data.energy[-1]
                if delta_en < packing_sim_data.energy[-2] / 10 and gtol > 2 * 1e-7:
                    gtol = gtol / np.sqrt(10)
                    print(f'Lowering gtol at {n_iter + 1}. minim minim for idx {idx}, now {gtol}')

            if energy > tol:
                raise ValueError('Check initialization, there should still be no overlaps for the first 10 iterations.')

        # increase (or decrease) the size of ellipses
        scaling += inflate_deflate * scaling_step
        ellipse_minim.change_params(scaling)

        minim_result = ellipse_minim.minim(eps=eps, gtol=gtol)
        packing_sim_data.append_all(minim_result, scaling, scaling_step)
        counter += 1
        n_iter += 1

    # save packing configuration and simulation data
    ellipse_minim.export_config(folder, idx=idx + existing_idx)
    packing_sim_data.export(folder, idx=idx + existing_idx)

    return scaling