コード例 #1
0
    def _divergence_clean(self, kx, ky, kz):

        mylog.info("Perform divergence cleaning.")

        self.gx = np.fft.fftn(self.gx)
        self.gy = np.fft.fftn(self.gy)
        self.gz = np.fft.fftn(self.gz)

        # These k's are different because we are
        # using the finite difference form of the
        # divergence operator.
        kxd = np.sin(kx * self.dx) / self.dx
        kyd = np.sin(ky * self.dy) / self.dy
        kzd = np.sin(kz * self.dz) / self.dz
        kkd = np.sqrt(kxd * kxd + kyd * kyd + kzd * kzd)
        with np.errstate(invalid='ignore', divide='ignore'):
            kxd /= kkd
            kyd /= kkd
            kzd /= kkd
            kxd[np.isnan(kxd)] = 0.0
            kyd[np.isnan(kyd)] = 0.0
            kzd[np.isnan(kzd)] = 0.0

        self.gx, self.gy, self.gz = \
            [(1.0-kxd*kxd)*self.gx-kxd*kyd*self.gy-kxd*kzd*self.gz,
             -kyd*kxd*self.gx+(1.0-kyd*kyd)*self.gy-kyd*kzd*self.gz,
             -kzd*kxd*self.gx-kzd*kyd*self.gy+(1.0-kzd*kzd)*self.gz]

        self.gx = np.fft.ifftn(self.gx).real
        self.gy = np.fft.ifftn(self.gy).real
        self.gz = np.fft.ifftn(self.gz).real

        del kxd, kyd, kzd, kkd
コード例 #2
0
ファイル: virial.py プロジェクト: jzuhone/cluster_generator
    def check_virial(self):
        r"""
        Computes the radial density profile for the collisionless 
        particles computed from integrating over the distribution 
        function, and the relative difference between this and the 
        input density profile.

        Returns
        -------
        rho : NumPy array
            The density profile computed from integrating the
            distribution function. 
        chk : NumPy array
            The relative difference between the input density
            profile and the one calculated using this method.
        """
        n = self.num_elements
        rho = np.zeros(n)
        pden = self.model[f"{self.ptype}_density"].d
        rho_int = lambda e, psi: self.f(e)*np.sqrt(2*(psi-e))
        for i, e in enumerate(self.ee):
            rho[i] = 4.*np.pi*quad(rho_int, 0., e, args=(e,))[0]
        chk = (rho[::-1]-pden)/pden
        mylog.info("The maximum relative deviation of this profile from "
                   "virial equilibrium is %g", np.abs(chk).max())
        return rho[::-1], chk
コード例 #3
0
ファイル: codes.py プロジェクト: jzuhone/cluster_generator
def setup_athena_ics(ics):
    r"""
    Parameters
    ----------
    ics : ClusterICs object
        The ClusterICs object to generate the Athena ICs from.
    """
    mylog.info("Add the following lines to athinput.cluster3d: ")
コード例 #4
0
    def from_dens_and_temp(cls,
                           rmin,
                           rmax,
                           density,
                           temperature,
                           stellar_density=None,
                           num_points=1000):
        """
        Construct a hydrostatic equilibrium model using gas density
        and temperature profiles. 

        Parameters
        ----------
        rmin : float
            Minimum radius of profiles in kpc.
        rmax : float
            Maximum radius of profiles in kpc.
        density : :class:`~cluster_generator.radial_profiles.RadialProfile`
            A radial profile describing the gas mass density.
        temperature : :class:`~cluster_generator.radial_profiles.RadialProfile`
            A radial profile describing the gas temperature.
        stellar_density : :class:`~cluster_generator.radial_profiles.RadialProfile`, optional
            A radial profile describing the stellar mass density, if desired.
        num_points : integer, optional
            The number of points the profiles are evaluated at.
        """
        mylog.info("Computing the profiles from density and temperature.")
        rr = np.logspace(np.log10(rmin),
                         np.log10(rmax),
                         num_points,
                         endpoint=True)
        fields = OrderedDict()
        fields["radius"] = unyt_array(rr, "kpc")
        fields["density"] = unyt_array(density(rr), "Msun/kpc**3")
        fields["temperature"] = unyt_array(temperature(rr), "keV")
        fields["pressure"] = fields["density"] * fields["temperature"]
        fields["pressure"] /= mu * mp
        fields["pressure"].convert_to_units("Msun/(Myr**2*kpc)")
        pressure_spline = InterpolatedUnivariateSpline(rr,
                                                       fields["pressure"].d)
        dPdr = unyt_array(pressure_spline(rr, 1), "Msun/(Myr**2*kpc**2)")
        fields["gravitational_field"] = dPdr / fields["density"]
        fields["gravitational_field"].convert_to_units("kpc/Myr**2")
        fields["gas_mass"] = unyt_array(integrate_mass(density, rr), "Msun")
        fields["total_mass"] = -fields["radius"]**2 * fields[
            "gravitational_field"] / G
        total_mass_spline = InterpolatedUnivariateSpline(
            rr, fields["total_mass"].v)
        dMdr = unyt_array(total_mass_spline(rr, nu=1), "Msun/kpc")
        fields["total_density"] = dMdr / (4. * np.pi * fields["radius"]**2)
        return cls._from_scratch(fields, stellar_density=stellar_density)
コード例 #5
0
    def _compute_vector_potential(self, kx, ky, kz):

        kk = np.sqrt(kx**2 + ky**2 + kz**2)

        mylog.info("Compute vector potential.")

        # Rotate vector potential

        self.gx = np.fft.fftn(self.gx)
        self.gy = np.fft.fftn(self.gy)
        self.gz = np.fft.fftn(self.gz)

        with np.errstate(invalid='ignore', divide='ignore'):
            alpha = np.arccos(kx / np.sqrt(kx * kx + ky * ky))
        alpha[ky < 0.0] -= 2.0 * np.pi
        alpha[ky < 0.0] *= -1.
        with np.errstate(invalid='ignore', divide='ignore'):
            beta = np.arccos(kz / kk)
        alpha[np.isinf(alpha)] = 0.0
        alpha[np.isnan(alpha)] = 0.0
        beta[np.isnan(beta)] = 0.0
        beta[np.isinf(beta)] = 0.0

        self._rot_3d(3, alpha)
        self._rot_3d(2, beta)

        with np.errstate(invalid='ignore', divide='ignore'):
            self.gx, self.gy = (
                0.0 + 1.0j) * self.gy / kk, -(0.0 + 1.0j) * self.gx / kk
            self.gz = np.zeros(self.gx.shape, dtype="complex")

        del kk

        self.gx[np.isinf(self.gx)] = 0.0
        self.gx[np.isnan(self.gx)] = 0.0
        self.gy[np.isinf(self.gy)] = 0.0
        self.gy[np.isnan(self.gy)] = 0.0

        self._rot_3d(2, -beta)
        self._rot_3d(3, -alpha)

        self.gx = np.fft.ifftn(self.gx).real
        self.gy = np.fft.ifftn(self.gy).real
        self.gz = np.fft.ifftn(self.gz).real
コード例 #6
0
ファイル: codes.py プロジェクト: jzuhone/cluster_generator
def setup_flash_ics(ics, use_particles=True, regenerate_particles=False):
    r"""

    Generate the "flash.par" lines needed for use
    with the GalaxyClusterMerger setup in FLASH. If the particles
    (dark matter and potentially star) have not been 
    created yet, they will be created at this step. 

    Parameters
    ----------
    ics : ClusterICs object
        The ClusterICs object to generate the GAMER ICs from.
    use_particles : boolean, optional
        If True, set up particle distributions. Default: True
    regenerate_particles : boolean, optional
        If particle files have already been created, particles
        are being used, and this flag is set to True, the particles 
        will be re-created. Default: False
    """
    if use_particles:
        ics._generate_particles(
            regenerate_particles=regenerate_particles)
    outlines = [
        f"testSingleCluster    {ics.num_halos} # number of halos"
    ]
    for i in range(ics.num_halos):
        vel = ics.velocity[i].to("km/s")
        outlines += [
            f"profile{i+1}\t=\t{ics.profiles[i]}\t# profile table of cluster {i+1}",
            f"xInit{i+1}\t=\t{ics.center[i][0]}\t# X-center of cluster {i+1} in kpc",
            f"yInit{i+1}\t=\t{ics.center[i][1]}\t# Y-center of cluster {i+1} in kpc",
            f"vxInit{i+1}\t=\t{vel[0]}\t# X-velocity of cluster {i+1} in km/s",
            f"vyInit{i+1}\t=\t{vel[1]}\t# Y-velocity of cluster {i+1} in km/s",
        ]
        if use_particles:
            outlines.append(
                f"Merger_File_Par{i+1}\t=\t{ics.particle_files[i]}\t# particle file of cluster {i+1}",
            )
    mylog.info("Add the following lines to flash.par: ")
    for line in outlines:
        print(line)
コード例 #7
0
    def no_gas(cls,
               rmin,
               rmax,
               total_density,
               stellar_density=None,
               num_points=1000):
        rr = np.logspace(np.log10(rmin),
                         np.log10(rmax),
                         num_points,
                         endpoint=True)
        fields = OrderedDict()
        fields["radius"] = unyt_array(rr, "kpc")
        fields["total_density"] = unyt_array(total_density(rr), "Msun/kpc**3")
        mylog.info("Integrating total mass profile.")
        fields["total_mass"] = unyt_array(integrate_mass(total_density, rr),
                                          "Msun")
        fields["gravitational_field"] = -G * fields["total_mass"] / (
            fields["radius"]**2)
        fields["gravitational_field"].convert_to_units("kpc/Myr**2")

        return cls._from_scratch(fields, stellar_density=stellar_density)
コード例 #8
0
    def check_hse(self):
        r"""
        Determine the deviation of the model from hydrostatic equilibrium.

        Returns
        -------
        chk : NumPy array
            An array containing the relative deviation from hydrostatic
            equilibrium as a function of radius.
        """
        if "pressure" not in self.fields:
            raise RuntimeError("This ClusterModel contains no gas!")
        rr = self.fields["radius"].v
        pressure_spline = InterpolatedUnivariateSpline(
            rr, self.fields["pressure"].v)
        dPdx = unyt_array(pressure_spline(rr, 1), "Msun/(Myr**2*kpc**2)")
        rhog = self.fields["density"] * self.fields["gravitational_field"]
        chk = dPdx - rhog
        chk /= rhog
        mylog.info(
            "The maximum relative deviation of this profile from "
            "hydrostatic equilibrium is %g",
            np.abs(chk).max())
        return chk
コード例 #9
0
    def from_dens_and_tden(cls,
                           rmin,
                           rmax,
                           density,
                           total_density,
                           stellar_density=None,
                           num_points=1000):
        """
        Construct a hydrostatic equilibrium model using gas density
        and total density profiles

        Parameters
        ----------
        rmin : float
            Minimum radius of profiles in kpc.
        rmax : float
            Maximum radius of profiles in kpc.
        density : :class:`~cluster_generator.radial_profiles.RadialProfile`
            A radial profile describing the gas mass density.
        total_density : :class:`~cluster_generator.radial_profiles.RadialProfile`
            A radial profile describing the total mass density.
        stellar_density : :class:`~cluster_generator.radial_profiles.RadialProfile`, optional
            A radial profile describing the stellar mass density, if desired.
        num_points : integer, optional
            The number of points the profiles are evaluated at.
        """
        mylog.info("Computing the profiles from density and total density.")
        rr = np.logspace(np.log10(rmin),
                         np.log10(rmax),
                         num_points,
                         endpoint=True)
        fields = OrderedDict()
        fields["radius"] = unyt_array(rr, "kpc")
        fields["density"] = unyt_array(density(rr), "Msun/kpc**3")
        fields["total_density"] = unyt_array(total_density(rr), "Msun/kpc**3")
        mylog.info("Integrating total mass profile.")
        fields["total_mass"] = unyt_array(integrate_mass(total_density, rr),
                                          "Msun")
        fields["gas_mass"] = unyt_array(integrate_mass(density, rr), "Msun")
        fields["gravitational_field"] = -G * fields["total_mass"] / (
            fields["radius"]**2)
        fields["gravitational_field"].convert_to_units("kpc/Myr**2")
        g = fields["gravitational_field"].in_units("kpc/Myr**2").v
        g_r = InterpolatedUnivariateSpline(rr, g)
        dPdr_int = lambda r: density(r) * g_r(r)
        mylog.info("Integrating pressure profile.")
        P = -integrate(dPdr_int, rr)
        dPdr_int2 = lambda r: density(r) * g[-1] * (rr[-1] / r)**2
        P -= quad(dPdr_int2, rr[-1], np.inf, limit=100)[0]
        fields["pressure"] = unyt_array(P, "Msun/kpc/Myr**2")
        fields[
            "temperature"] = fields["pressure"] * mu * mp / fields["density"]
        fields["temperature"].convert_to_units("keV")

        return cls._from_scratch(fields, stellar_density=stellar_density)
コード例 #10
0
    def _from_scratch(cls, fields, stellar_density=None):
        rr = fields["radius"].d
        mylog.info("Integrating gravitational potential profile.")
        tdens_func = InterpolatedUnivariateSpline(rr,
                                                  fields["total_density"].d)
        gpot_profile = lambda r: tdens_func(r) * r
        gpot1 = fields["total_mass"] / fields["radius"]
        gpot2 = unyt_array(4. * np.pi * integrate(gpot_profile, rr),
                           "Msun/kpc")
        fields["gravitational_potential"] = -G * (gpot1 + gpot2)
        fields["gravitational_potential"].convert_to_units("kpc**2/Myr**2")

        if "density" in fields and "gas_mass" not in fields:
            mylog.info("Integrating gas mass profile.")
            m0 = fields["density"].d[0] * rr[0]**3 / 3.
            fields["gas_mass"] = unyt_array(
                4.0 * np.pi *
                cumtrapz(fields["density"] * rr * rr, x=rr, initial=0.0) + m0,
                "Msun")

        if stellar_density is not None:
            fields["stellar_density"] = unyt_array(stellar_density(rr),
                                                   "Msun/kpc**3")
            mylog.info("Integrating stellar mass profile.")
            fields["stellar_mass"] = unyt_array(
                integrate_mass(stellar_density, rr), "Msun")

        mdm = fields["total_mass"].copy()
        ddm = fields["total_density"].copy()
        if "density" in fields:
            mdm -= fields["gas_mass"]
            ddm -= fields["density"]
        if "stellar_mass" in fields:
            mdm -= fields["stellar_mass"]
            ddm -= fields["stellar_density"]
        mdm[ddm.v < 0.0] = mdm.max()
        ddm[ddm.v < 0.0] = 0.0

        if ddm.sum() < 0.0 or mdm.sum() < 0.0:
            mylog.warning(
                "The total dark matter mass is either zero or negative!!")
        fields["dark_matter_density"] = ddm
        fields["dark_matter_mass"] = mdm

        if "density" in fields:
            fields["gas_fraction"] = fields["gas_mass"] / fields["total_mass"]
            fields["electron_number_density"] = \
                fields["density"].to("cm**-3", "number_density", mu=mue)
            fields["entropy"] = \
                fields["temperature"]*fields["electron_number_density"]**mtt

        return cls(rr.size, fields)
コード例 #11
0
    def generate_gas_particles(self,
                               num_particles,
                               r_max=None,
                               sub_sample=1,
                               compute_potential=False,
                               prng=None):
        """
        Generate a set of gas particles in hydrostatic equilibrium.

        Parameters
        ----------
        num_particles : integer
            The number of particles to generate.
        r_max : float, optional
            The maximum radius in kpc within which to generate 
            particle positions. If not supplied, it will generate
            positions out to the maximum radius available. Default: None
        sub_sample : integer, optional
            This option allows one to generate a sub-sample of unique
            particle radii, densities, and energies which will then be
            repeated to fill the required number of particles. Default: 1,
            which means no sub-sampling.
        compute_potential : boolean, optional
            If True, the gravitational potential for each particle will
            be computed. Default: False
        prng : :class:`~numpy.random.RandomState` object, integer, or None
            A pseudo-random number generator. Typically will only 
            be specified if you have a reason to generate the same 
            set of random numbers, such as for a test. Default is None, 
            which sets the seed based on the system time.
        """
        from cluster_generator.utils import parse_prng
        prng = parse_prng(prng)
        mylog.info("We will be assigning %d particles." % num_particles)
        mylog.info("Compute particle positions.")

        num_particles_sub = num_particles // sub_sample

        radius_sub, mtot = generate_particle_radii(self["radius"].d,
                                                   self["gas_mass"].d,
                                                   num_particles_sub,
                                                   r_max=r_max,
                                                   prng=prng)

        if sub_sample > 1:
            radius = np.tile(radius_sub, sub_sample)[:num_particles]
        else:
            radius = radius_sub

        theta = np.arccos(prng.uniform(low=-1., high=1., size=num_particles))
        phi = 2. * np.pi * prng.uniform(size=num_particles)

        fields = OrderedDict()

        fields["gas", "particle_position"] = unyt_array([
            radius * np.sin(theta) * np.cos(phi),
            radius * np.sin(theta) * np.sin(phi), radius * np.cos(theta)
        ], "kpc").T

        mylog.info("Compute particle thermal energies, densities, and masses.")

        e_arr = 1.5 * self.fields["pressure"] / self.fields["density"]
        get_energy = InterpolatedUnivariateSpline(self.fields["radius"], e_arr)

        if sub_sample > 1:
            energy = np.tile(get_energy(radius_sub),
                             sub_sample)[:num_particles]
        else:
            energy = get_energy(radius)

        fields["gas", "thermal_energy"] = unyt_array(energy, "kpc**2/Myr**2")
        fields["gas", "particle_mass"] = unyt_array([mtot / num_particles] *
                                                    num_particles, "Msun")

        get_density = InterpolatedUnivariateSpline(self.fields["radius"],
                                                   self.fields["density"])

        if sub_sample > 1:
            density = np.tile(get_density(radius_sub),
                              sub_sample)[:num_particles]
        else:
            density = get_density(radius)

        fields["gas", "density"] = unyt_array(density, "Msun/kpc**3")

        mylog.info("Set particle velocities to zero.")

        fields["gas",
               "particle_velocity"] = unyt_array(np.zeros((num_particles, 3)),
                                                 "kpc/Myr")

        if compute_potential:
            energy_spline = InterpolatedUnivariateSpline(
                self["radius"].d, -self["gravitational_potential"])
            phi = -energy_spline(radius_sub)
            if sub_sample > 1:
                phi = np.tile(phi, sub_sample)
            fields["gas",
                   "particle_potential"] = unyt_array(phi, "kpc**2/Myr**2")

        return ClusterParticles("gas", fields)
コード例 #12
0
ファイル: codes.py プロジェクト: jzuhone/cluster_generator
def setup_gamer_ics(ics, regenerate_particles=False, use_tracers=False):
    r"""

    Generate the "Input_TestProb" lines needed for use
    with the ClusterMerger setup in GAMER. If the particles
    (dark matter and potentially star) have not been 
    created yet, they will be created at this step. New profile 
    files will also be created which have all fields in CGS units
    for reading into GAMER. If a magnetic field file is present
    in the ICs, a note will be given about how it should be named
    for GAMER to use it.

    NOTE: Gas particles in the initial conditions will be interpreted
    as tracer particles.

    Parameters
    ----------
    ics : ClusterICs object
        The ClusterICs object to generate the GAMER ICs from.
    regenerate_particles : boolean, optional
        If particle files have already been created and this
        flag is set to True, the particles will be
        re-created. Default: False
    use_tracers : boolean
        Set to True to add tracer particles. Default: False
    """
    gamer_ptypes = ["dm", "star"]
    if use_tracers:
        gamer_ptypes.insert(0, "gas")
    gamer_ptype_num = {"gas": 0, "dm": 2, "star": 3}
    hses = [ClusterModel.from_h5_file(hf) for hf in ics.profiles]
    parts = ics._generate_particles(
        regenerate_particles=regenerate_particles)
    outlines = [
        f"Merger_Coll_NumHalos\t\t{ics.num_halos}\t# number of halos"
    ]
    for i in range(ics.num_halos):
        particle_file = f"{ics.basename}_gamerp_{i+1}.h5"
        parts[i].write_sim_input(particle_file, gamer_ptypes, gamer_ptype_num)
        hse_file_gamer = ics.profiles[i].replace(".h5", "_gamer.h5")
        hses[i].write_model_to_h5(hse_file_gamer, overwrite=True,
                                  in_cgs=True, r_max=ics.r_max)
        vel = ics.velocity[i].to_value("km/s")
        outlines += [
            f"Merger_File_Prof{i+1}\t\t{hse_file_gamer}\t# profile table of cluster {i+1}",
            f"Merger_File_Par{i+1}\t\t{particle_file}\t# particle file of cluster {i+1}",
            f"Merger_Coll_PosX{i+1}\t\t{ics.center[i][0].v}\t# X-center of cluster {i+1} in kpc",
            f"Merger_Coll_PosY{i+1}\t\t{ics.center[i][1].v}\t# Y-center of cluster {i+1} in kpc",
            f"Merger_Coll_VelX{i+1}\t\t{vel[0]}\t# X-velocity of cluster {i+1} in km/s",
            f"Merger_Coll_VelY{i+1}\t\t{vel[1]}\t# Y-velocity of cluster {i+1} in km/s"
        ]
    mylog.info("Write the following lines to Input__TestProblem: ")
    for line in outlines:
        print(line)
    num_particles = sum([ics.tot_np[key] for key in ics.tot_np])
    mylog.info(f"In the Input__Parameter file, "
               f"set PAR__NPAR = {num_particles}.")
    if ics.mag_file is not None:
        mylog.info(f"Rename the file '{ics.mag_file}' to 'B_IC' "
                   f"and place it in the same directory as the "
                   f"Input__* files, and set OPT__INIT_BFIELD_BYFILE "
                   f"to 1 in Input__Parameter")
コード例 #13
0
    def __init__(self,
                 left_edge,
                 right_edge,
                 ddims,
                 l_min,
                 l_max,
                 padding=0.1,
                 alpha=-11. / 3.,
                 g_rms=1.0,
                 ctr1=None,
                 ctr2=None,
                 ctr3=None,
                 r1=None,
                 r2=None,
                 r3=None,
                 g1=None,
                 g2=None,
                 g3=None,
                 vector_potential=False,
                 divergence_clean=False,
                 prng=None,
                 r_max=None):

        prng = parse_prng(prng)

        super(GaussianRandomField,
              self).__init__(left_edge,
                             right_edge,
                             ddims,
                             padding=padding,
                             vector_potential=vector_potential,
                             divergence_clean=divergence_clean)

        nx, ny, nz = self.ddims

        num_halos = 0
        if r1 is not None:
            num_halos += 1
        if r2 is not None:
            num_halos += 1
        if r3 is not None:
            num_halos += 1

        if num_halos >= 1:
            if ctr1 is None:
                ctr1 = 0.5 * (self.left_edge + self.right_edge)
            else:
                ctr1 = parse_value(ctr1, "kpc").v
            r1 = parse_value(r1, "kpc").v
            g1 = parse_value(g1, self._units)
        if num_halos >= 2:
            if ctr2 is None:
                raise RuntimeError(
                    "Need to specify 'ctr2' for the second halo!")
            ctr2 = parse_value(ctr2, "kpc").v
            r2 = parse_value(r2, "kpc").v
            g2 = parse_value(g2, self._units)
        if num_halos == 3:
            if ctr3 is None:
                raise RuntimeError(
                    "Need to specify 'ctr3' for the second halo!")
            ctr3 = parse_value(ctr3, "kpc").v
            r3 = parse_value(r3, "kpc").v
            g3 = parse_value(g3, self._units)

        # Derived stuff

        l_min = parse_value(l_min, "kpc").v
        l_max = parse_value(l_max, "kpc").v

        k0 = 2. * np.pi / l_min
        k1 = 2. * np.pi / l_max

        mylog.info("Setting up the Gaussian random fields.")

        v = np.exp(2. * np.pi * 1j * prng.random((3, nx, ny, nz)))

        v[:, 0, 0, 0] = 2. * np.sign(
            (v[:, 0, 0, 0].imag < 0.0).astype("int")) - 1. + 0j
        v[:, nx // 2, ny // 2, nz // 2] = 2. * np.sign(
            (v[:, nx // 2, ny // 2, nz // 2].imag <
             0.0).astype("int")) - 1. + 0j
        v[:, 0, ny // 2, nz // 2] = 2. * np.sign(
            (v[:, 0, ny // 2, nz // 2].imag < 0.0).astype("int")) - 1. + 0j
        v[:, nx // 2, 0, nz // 2] = 2. * np.sign(
            (v[:, nx // 2, 0, nz // 2].imag < 0.0).astype("int")) - 1. + 0j
        v[:, nx // 2, ny // 2, 0] = 2. * np.sign(
            (v[:, nx // 2, ny // 2, 0].imag < np.pi).astype("int")) - 1. + 0j
        v[:, 0, 0, nz // 2] = 2. * np.sign(
            (v[:, 0, 0, nz // 2].imag < np.pi).astype("int")) - 1. + 0j
        v[:, 0, ny // 2, 0] = 2. * np.sign(
            (v[:, 0, ny // 2, 0].imag < np.pi).astype("int")) - 1. + 0j
        v[:, nx // 2, 0, 0] = 2. * np.sign(
            (v[:, nx // 2, 0, 0].imag < np.pi).astype("int")) - 1. + 0j

        np.multiply(v, np.sqrt(-2. * np.log(prng.random((3, nx, ny, nz)))), v)

        kx, ky, kz = self._compute_waves()
        kk = np.sqrt(kx**2 + ky**2 + kz**2)
        with np.errstate(invalid='ignore', divide='ignore'):
            sigma = (1.0 + (kk / k1)**2)**(0.25 * alpha) * np.exp(-0.5 *
                                                                  (kk / k0)**2)
        sigma[np.isinf(sigma)] = 0.0
        sigma[np.isnan(sigma)] = 0.0
        del kk

        v[:, nx - 1:0:-1, ny - 1:0:-1,
          nz - 1:nz // 2:-1] = np.conjugate(v[:, 1:nx, 1:ny, 1:nz // 2])
        v[:, nx - 1:0:-1, ny - 1:ny // 2:-1,
          nz // 2] = np.conjugate(v[:, 1:nx, 1:ny // 2, nz // 2])
        v[:, nx - 1:0:-1, ny - 1:ny // 2:-1, 0] = np.conjugate(v[:, 1:nx,
                                                                 1:ny // 2, 0])
        v[:, nx - 1:0:-1, 0, nz - 1:nz // 2:-1] = np.conjugate(v[:, 1:nx, 0,
                                                                 1:nz // 2])
        v[:, 0, ny - 1:0:-1, nz - 1:nz // 2:-1] = np.conjugate(v[:, 0, 1:ny,
                                                                 1:nz // 2])
        v[:, nx - 1:nx // 2:-1, ny // 2,
          nz // 2] = np.conjugate(v[:, 1:nx // 2, ny // 2, nz // 2])
        v[:, nx - 1:nx // 2:-1, ny // 2, 0] = np.conjugate(v[:, 1:nx // 2,
                                                             ny // 2, 0])
        v[:, nx - 1:nx // 2:-1, 0, nz // 2] = np.conjugate(v[:, 1:nx // 2, 0,
                                                             nz // 2])
        v[:, 0, ny - 1:ny // 2:-1, nz // 2] = np.conjugate(v[:, 0, 1:ny // 2,
                                                             nz // 2])
        v[:, nx - 1:nx // 2:-1, 0, 0] = np.conjugate(v[:, 1:nx // 2, 0, 0])
        v[:, 0, ny - 1:ny // 2:-1, 0] = np.conjugate(v[:, 0, 1:ny // 2, 0])
        v[:, 0, 0, nz - 1:nz // 2:-1] = np.conjugate(v[:, 0, 0, 1:nz // 2])

        gx = np.fft.ifftn(sigma * v[0, :, :, :]).real
        gy = np.fft.ifftn(sigma * v[1, :, :, :]).real
        gz = np.fft.ifftn(sigma * v[2, :, :, :]).real

        del sigma, v

        g_avg = np.sqrt(np.mean(gx * gx + gy * gy + gz * gz))

        gx /= g_avg
        gy /= g_avg
        gz /= g_avg

        x, y, z = self._compute_coords()

        if num_halos == 0:
            g_rms = parse_value(g_rms, self._units)
            mylog.info(f"Scaling the fields by the constant value {g_rms}.")
        else:
            if num_halos >= 1:
                mylog.info("Scaling the fields by cluster 1.")
                rr1 = np.sqrt((x - ctr1[0])**2 + (y - ctr1[1])**2 +
                              (z - ctr1[2])**2)
                if r_max is not None:
                    rr1[rr1 > r_max] = r_max
                idxs1 = np.searchsorted(r1, rr1) - 1
                dr1 = (rr1 - r1[idxs1]) / (r1[idxs1 + 1] - r1[idxs1])
                g_rms = ((1. - dr1) * g1[idxs1] + dr1 * g1[idxs1 + 1])**2
            if num_halos >= 2:
                mylog.info("Scaling the fields by cluster 2.")
                rr2 = np.sqrt((x - ctr2[0])**2 + (y - ctr2[1])**2 +
                              (z - ctr2[2])**2)
                if r_max is not None:
                    rr2[rr2 > r_max] = r_max
                idxs2 = np.searchsorted(r2, rr2) - 1
                dr2 = (rr2 - r2[idxs2]) / (r2[idxs2 + 1] - r2[idxs2])
                g_rms += ((1. - dr2) * g2[idxs2] + dr2 * g2[idxs2 + 1])**2
            if num_halos == 3:
                mylog.info("Scaling the fields by cluster 3.")
                rr3 = np.sqrt((x - ctr3[0])**2 + (y - ctr3[1])**2 +
                              (z - ctr3[2])**2)
                if r_max is not None:
                    rr3[rr3 > r_max] = r_max
                idxs3 = np.searchsorted(r3, rr3) - 1
                dr3 = (rr3 - r3[idxs3]) / (r3[idxs3 + 1] - r3[idxs3])
                g_rms += ((1. - dr3) * g3[idxs3] + dr3 * g3[idxs3 + 1])**2
            g_rms = np.sqrt(g_rms).in_units(self._units).d

        gx *= g_rms
        gy *= g_rms
        gz *= g_rms

        self.x = x[:, 0, 0]
        self.y = y[0, :, 0]
        self.z = z[0, 0, :]

        self.gx = gx
        self.gy = gy
        self.gz = gz

        rescale = (self.gx**2 + self.gy**2 + self.gz**2).sum()

        del x, y, z, g_rms

        if self.divergence_clean:
            self._divergence_clean(kx, ky, kz)

        rescale /= (self.gx**2 + self.gy**2 + self.gz**2).sum()

        self.gx *= rescale
        self.gy *= rescale
        self.gz *= rescale

        if self.vector_potential:
            self._compute_vector_potential(kx, ky, kz)
コード例 #14
0
ファイル: virial.py プロジェクト: jzuhone/cluster_generator
    def generate_particles(self, num_particles, r_max=None, sub_sample=1,
                           compute_potential=False, prng=None):
        """
        Generate a set of dark matter or star particles in virial equilibrium.

        Parameters
        ----------
        num_particles : integer
            The number of particles to generate.
        r_max : float, optional
            The maximum radius in kpc within which to generate 
            particle positions. If not supplied, it will generate
            positions out to the maximum radius available. Default: None
        sub_sample : integer, optional
            This option allows one to generate a sub-sample of unique
            particle radii and velocities which will then be repeated
            to fill the required number of particles. Default: 1, which
            means no sub-sampling.
        compute_potential : boolean, optional
            If True, the gravitational potential for each particle will
            be computed. Default: False
        prng : :class:`~numpy.random.RandomState` object, integer, or None
            A pseudo-random number generator. Typically will only 
            be specified if you have a reason to generate the same 
            set of random numbers, such as for a test. Default is None, 
            which sets the seed based on the system time.

        Returns
        -------
        particles : :class:`~cluster_generator.particles.ClusterParticles`
            A set of dark matter or star particles.
        """
        from cluster_generator.utils import parse_prng

        num_particles_sub = num_particles // sub_sample
        key = {"dark_matter": "dm",  "stellar": "star"}[self.ptype]
        density = f"{self.ptype}_density"
        mass = f"{self.ptype}_mass"
        energy_spline = InterpolatedUnivariateSpline(self.model["radius"].d,
                                                     self.ee[::-1])

        prng = parse_prng(prng)

        mylog.info("We will be assigning %s %s particles.",
                   num_particles, self.ptype)
        mylog.info("Compute %s particle positions.", num_particles)

        nonzero = self.model[density] > 0.0
        radius_sub, mtot = generate_particle_radii(self.model["radius"].d[nonzero],
                                                   self.model[mass].d[nonzero],
                                                   num_particles_sub, r_max=r_max,
                                                   prng=prng)

        if sub_sample > 1:
            radius = np.tile(radius_sub, sub_sample)[:num_particles]
        else:
            radius = radius_sub

        theta = np.arccos(prng.uniform(low=-1., high=1., size=num_particles))
        phi = 2.*np.pi*prng.uniform(size=num_particles)

        fields = OrderedDict()

        fields[key, "particle_position"] = unyt_array(
            [radius*np.sin(theta)*np.cos(phi), radius*np.sin(theta)*np.sin(phi),
             radius*np.cos(theta)], "kpc").T

        mylog.info("Compute %s particle velocities.", self.ptype)

        psi = energy_spline(radius_sub)
        vesc = 2.*psi
        fv2esc = vesc*self.f(psi)
        vesc = np.sqrt(vesc)

        velocity_sub = generate_velocities(
            psi, vesc, fv2esc, self.f._eval_args[0], self.f._eval_args[1],
            self.f._eval_args[2])

        if sub_sample > 1:
            velocity = np.tile(velocity_sub, sub_sample)[:num_particles]
        else:
            velocity = velocity_sub

        theta = np.arccos(prng.uniform(low=-1., high=1., size=num_particles))
        phi = 2.*np.pi*prng.uniform(size=num_particles)

        fields[key, "particle_velocity"] = unyt_array(
            [velocity*np.sin(theta)*np.cos(phi),
             velocity*np.sin(theta)*np.sin(phi),
             velocity*np.cos(theta)], "kpc/Myr").T

        fields[key, "particle_mass"] = unyt_array(
            [mtot/num_particles]*num_particles, "Msun")

        if compute_potential:
            if sub_sample > 1:
                phi = -np.tile(psi, sub_sample)
            else:
                phi = -psi
            fields[key, "particle_potential"] = unyt_array(
                phi, "kpc**2/Myr**2")

        return ClusterParticles(key, fields)