AMSS-NCKU/puncture_initialize.py

##################################################################
##
## Initialize puncture data used by AMSS-NCKU Python helpers
## Author: Xiaoqu
## Adapted for the current branch
##
##################################################################

import math

import numpy

import AMSS_NCKU_Input as input_data


class PunctureData:
    def __init__(
        self,
        BBH_M1,
        BBH_M2,
        distance_d0,
        ellipticity_e0,
        dimensionless_mass_BH,
        charge_Q_BH,
        position_BH,
        momentum_BH,
        angular_momentum_BH,
    ):
        self.BBH_M1 = BBH_M1
        self.BBH_M2 = BBH_M2
        self.distance_d0 = distance_d0
        self.ellipticity_e0 = ellipticity_e0
        self.dimensionless_mass_BH = dimensionless_mass_BH
        self.charge_Q_BH = charge_Q_BH
        self.position_BH = position_BH
        self.momentum_BH = momentum_BH
        self.angular_momentum_BH = angular_momentum_BH


def _angular_momentum_from_input(masses):
    angular_momentum_BH = numpy.zeros((input_data.puncture_number, 3))
    for i in range(input_data.puncture_number):
        if input_data.Symmetry in ("octant-symmetry", "equatorial-symmetry"):
            angular_momentum_BH[i] = [0.0, 0.0, (masses[i] ** 2) * input_data.parameter_BH[i, 2]]
        elif input_data.Symmetry == "no-symmetry":
            angular_momentum_BH[i] = (masses[i] ** 2) * input_data.dimensionless_spin_BH[i]
        else:
            raise ValueError("Unsupported symmetry setting")
    return angular_momentum_BH


def print_puncture_information(puncture_data):
    print("------------------------------------------------------------------------------------------")
    print()
    print(" Printing the puncture information ")
    print()

    for i in range(input_data.puncture_number):
        mass = puncture_data.dimensionless_mass_BH[i]
        charge = puncture_data.charge_Q_BH[i]
        position = puncture_data.position_BH[i]
        momentum = puncture_data.momentum_BH[i]
        angular_momentum = puncture_data.angular_momentum_BH[i]

        if input_data.Symmetry in ("octant-symmetry", "equatorial-symmetry"):
            a_star = angular_momentum[2] / (mass ** 2)
        else:
            a_star = numpy.linalg.norm(angular_momentum) / (mass ** 2)

        print(f" The information for puncture {i+1} ")
        print(f" Mass({i+1}) = {mass:>10.6f},  Q({i+1})  = {charge:>10.6f},  a*({i+1}) = {a_star:>10.6f}")
        print(f" X({i+1})    = {position[0]:>10.6f},  Y({i+1})  = {position[1]:>10.6f},  Z({i+1})  = {position[2]:>10.6f}")
        print(f" Px({i+1})   = {momentum[0]:>10.6f},  Py({i+1}) = {momentum[1]:>10.6f},  Pz({i+1}) = {momentum[2]:>10.6f}")
        print(
            f" Jx({i+1})   = {angular_momentum[0]:>10.6f},  Jy({i+1}) = {angular_momentum[1]:>10.6f},  Jz({i+1}) = {angular_momentum[2]:>10.6f}"
        )
        print()

    print("------------------------------------------------------------------------------------------")


def generate_puncture_input_data():
    print()
    print(" Setting puncture position, momentum and angular momentum ")
    print()

    if input_data.puncture_data_set == "Automatically-BBH":
        mass_ratio_Q = input_data.parameter_BH[0, 0] / input_data.parameter_BH[1, 0]
        if mass_ratio_Q < 1.0:
            raise ValueError("The first black hole must be the larger-mass puncture")

        BBH_M1 = mass_ratio_Q / (1.0 + mass_ratio_Q)
        BBH_M2 = 1.0 / (1.0 + mass_ratio_Q)
        distance_d0 = input_data.Distance
        ellipticity_e0 = input_data.e0

        position_BH = numpy.zeros((input_data.puncture_number, 3))
        position_BH[0] = [0.0, distance_d0 / (1.0 + mass_ratio_Q), 0.0]
        position_BH[1] = [0.0, -distance_d0 * mass_ratio_Q / (1.0 + mass_ratio_Q), 0.0]
        for i in range(2, input_data.puncture_number):
            position_BH[i] = input_data.position_BH[i]

        dimensionless_mass_BH = numpy.zeros(input_data.puncture_number)
        dimensionless_mass_BH[0] = BBH_M1
        dimensionless_mass_BH[1] = BBH_M2
        for i in range(2, input_data.puncture_number):
            dimensionless_mass_BH[i] = input_data.parameter_BH[i, 0]

        charge_Q_BH = dimensionless_mass_BH * input_data.parameter_BH[:, 1]
        angular_momentum_BH = _angular_momentum_from_input(dimensionless_mass_BH)

        import BBH_orbit_parameter

        BBH_S1 = angular_momentum_BH[0] / (BBH_M1 ** 2)
        BBH_S2 = angular_momentum_BH[1] / (BBH_M2 ** 2)
        momentum_BH = numpy.zeros((input_data.puncture_number, 3))
        momentum_BH[0], momentum_BH[1] = BBH_orbit_parameter.generate_BBH_orbit_parameters(
            BBH_M1, BBH_M2, BBH_S1, BBH_S2, distance_d0, ellipticity_e0
        )
        for i in range(2, input_data.puncture_number):
            momentum_BH[i] = input_data.momentum_BH[i]

    elif input_data.puncture_data_set == "Manually":
        position_BH = numpy.array(input_data.position_BH, copy=True)
        momentum_BH = numpy.array(input_data.momentum_BH, copy=True)
        dimensionless_mass_BH = numpy.array(input_data.parameter_BH[:, 0], copy=True)
        charge_Q_BH = dimensionless_mass_BH * input_data.parameter_BH[:, 1]
        angular_momentum_BH = _angular_momentum_from_input(dimensionless_mass_BH)

        if input_data.puncture_number >= 2:
            distance_d0 = math.sqrt(
                (position_BH[0, 0] - position_BH[1, 0]) ** 2
                + (position_BH[0, 1] - position_BH[1, 1]) ** 2
                + (position_BH[0, 2] - position_BH[1, 2]) ** 2
            )
        else:
            distance_d0 = 0.0
        ellipticity_e0 = input_data.e0
        BBH_M1 = dimensionless_mass_BH[0] if input_data.puncture_number >= 1 else 0.0
        BBH_M2 = dimensionless_mass_BH[1] if input_data.puncture_number >= 2 else 0.0

    else:
        raise ValueError("Unsupported puncture_data_set setting")

    puncture_data = PunctureData(
        BBH_M1,
        BBH_M2,
        distance_d0,
        ellipticity_e0,
        dimensionless_mass_BH,
        charge_Q_BH,
        position_BH,
        momentum_BH,
        angular_momentum_BH,
    )
    print_puncture_information(puncture_data)
    return puncture_data