<commit_before># -*- coding: utf-8 -*-
"""
Contains functions for genrerating a simulation snapshot from an IC object 
which already has particle positions generated.  Snapshot are made with
snapshot_gen(IC)

Created on Fri Mar 21 15:11:31 2014

@author: ibackus
"""

__version__ = "$Revision: 2 $"
# $Source$

import pynbody
SimArray = pynbody.array.SimArray
import gc

# diskpy packages
import diskpy
from diskpy import global_settings
from diskpy.utils import match_units, strip_units
from diskpy.pychanga import make_director, make_param, setup_units
import calc_velocity

# Constants
G = SimArray(1.0,'G')

def snapshot_gen(IC):
    """
    Generates a tipsy snapshot from the initial conditions object IC.  Includes
    a routine to calculate the velocity.
    
    Parameters
    ----------
    IC : ICobj
    
    Returns
    -------
    snapshot : SimSnap
        Simulation snapshot with the velocity calculated
    param : dict
        dictionary containing info for a .param file
    """
    
    print 'Generating snapshot...'   
    # Initialize snapshot
    snapshot = init_snapshot(IC)
    # Make param file
    param = init_param(IC, snapshot)
       
    # -------------------------------------------------
    # CALCULATE VELOCITY USING calc_velocity.py.  This also estimates the 
    # gravitational softening length eps and a good timestep
    # -------------------------------------------------
    print 'Calculating circular velocity'
    preset = IC.settings.changa_run.preset
    max_particles = global_settings['misc']['max_particles']
    dDelta = calc_velocity.v_xy(snapshot, param, \
    changa_preset=preset, max_particles=max_particles)
    param['dDelta'] = dDelta    
    print 'Calculated time step.  dDelta = ', dDelta
    gc.collect()
    
    # Create director file
    director = init_director(IC, param)
    # Finalize
    print 'Wrapping up'
    setup_sinks(snapshot, param, r_sink = IC.pos.r.min())
    
    return snapshot, param, director

def init_director(IC, param=None):
    """
    Creates a director dict for a PPD.
    
    Parameters
    ----------
    IC : ICobj
    param : .param dict
        If not supplied, param = IC.param
    """    
    if param is None:
        
        param = IC.snapshot_param
        
    # largest radius to plot
    r_director = float(0.9 * IC.pos.r.max())
    # Maximum surface density
    sigma_min = float(IC.sigma(r_director))
    # surface density at largest radius
    sigma_max = float(IC.sigma.input_dict['sigma'].max())
    # Create director dict
    director = make_director(sigma_min, sigma_max, r_director, \
    filename=param['achOutName'])
    
    return director
    
def setup_sinks(snapshot, param, r_sink=None):
    """
    Sets up snapshot and param for stars that are sinks.
    
    Parameters
    ----------
    snapshot : SimSnap
    param : param dict
    r_sink : SimArray-like
        If provide, this will be the sink radius.  Otherwise it's taken to be
        the minimum gas cylindrical radius
    
    Returns
    -------
    None
    """
    units = diskpy.pychanga.units_from_param(param)
    
    if r_sink is None:
        
        r_sink = snapshot.g['rxy'].min()
        r_sink.convert_units(units['l_unit'])
        
    r_sink = strip_units(r_sink)
    # Set the star tforms to a negative number.  This allows UW ChaNGa treat 
    # stars as sink particles
    snapshot.star['tform'] = -1.0    
    # Update params
    Mstar = snapshot.s['mass'].min()
    Mstar.convert_units(units['m_unit'])
    param['dSinkBoundOrbitRadius'] = r_sink
    param['dSinkRadius'] = r_sink
    param['dSinkMassMin'] = 0.9 * strip_units(Mstar)
    param['bDoSinks'] = 1

def init_param(IC, snapshot=None):
    """
    Initializes a ChaNGa param dict (see also diskpy.pychanga.make_param) for
    an IC object.
    
    Parameters
    ----------
    IC : ICobject
        Initial conditions object containing required settings for generating
        the param dict
    snapshot : SimSnap<commit_msg>Completed porting make_snapshot_binary over to make_snapshot<commit_after># -*- coding: utf-8 -*-
"""
Contains functions for genrerating a simulation snapshot from an IC object 
which already has particle positions generated.  Snapshot are made with
snapshot_gen(IC)

Created on Fri Mar 21 15:11:31 2014

@author: ibackus
"""

__version__ = "$Revision: 2 $"
# $Source$

import pynbody
SimArray = pynbody.array.SimArray
import gc
import numpy as np

# diskpy packages
import diskpy
from diskpy import global_settings
from diskpy.utils import match_units, strip_units
from diskpy.pychanga import make_director, make_param, setup_units
import calc_velocity
import AddBinary

# Constants
G = SimArray(1.0,'G')

def snapshot_gen(IC):
    """
    Generates a tipsy snapshot from the initial conditions object IC.  Includes
    a routine to calculate the velocity.
    
    Parameters
    ----------
    IC : ICobj
    
    Returns
    -------
    snapshot : SimSnap
        Simulation snapshot with the velocity calculated
    param : dict
        dictionary containing info for a .param file
    """
    
    print 'Generating snapshot...'   
    # Initialize snapshot
    snapshot = init_snapshot(IC)
    # Make param file
    param = init_param(IC, snapshot)
       
    # -------------------------------------------------
    # CALCULATE VELOCITY USING calc_velocity.py.  This also estimates the 
    # gravitational softening length eps and a good timestep
    # -------------------------------------------------
    print 'Calculating circular velocity'
    preset = IC.settings.changa_run.preset
    max_particles = global_settings['misc']['max_particles']
    dDelta = calc_velocity.v_xy(snapshot, param, \
    changa_preset=preset, max_particles=max_particles)
    param['dDelta'] = dDelta    
    print 'Calculated time step.  dDelta = ', dDelta
    gc.collect()
    
    # Create director file
    director = init_director(IC, param)
    # Handle binaries
    starMode = IC.settings.physical.starMode.lower()
    if starMode is 'binary':
        
        snapshot = make_binary(IC, snapshot)
        
    # Finalize
    print 'Wrapping up'
    setup_sinks(IC, snapshot, param)
    
    return snapshot, param, director
    
def make_binary(IC, snapshot):
    """
    Turns a snapshot for a single star into a snapshot of a binary system
    
    Parameters
    ----------
    IC : IC object
    snapshot : SimSnap
        Single star system to turn into a binary
    
    Returns
    -------
    snapshotBinary : SimSnap
        A binary version of the simulation snapshot
    """
    snapshotBinary = pynbody.new(star=2, gas=len(snapshot.g))
    # Copy gas particles over
    for key in snapshot.gas.keys():
        
        snapshotBinary.gas[key] = snapshot.gas[key]
        
    # Load Binary system obj to initialize system
    binsys = IC.settings.physical.binsys
    x1,x2,v1,v2 = binsys.generateICs()
    
    #Assign star parameters assuming CCW orbit
    snapshotBinary.star[0]['pos'] = x1
    snapshotBinary.star[0]['vel'] = v1
    snapshotBinary.star[1]['pos'] = x2
    snapshotBinary.star[1]['vel'] = v2
    
    #Set stellar masses
    priMass = binsys.m1
    secMass = binsys.m2    
    snapshotBinary.star[0]['mass'] = priMass
    snapshotBinary.star[1]['mass'] = secMass
    snapshotBinary.star['metals'] = snapshot.s['metals']
    
    # Estimate stars' softening length as fraction of distance to COM
    d = np.sqrt(AddBinary.dotProduct(x1-x2,x1-x2))
    pos_unit = snapshotBinary['pos'].units
    snapshotBinary.star['eps'] = SimArray(abs(d)/4.0,pos_unit)
    
    return snapshotBinary

def init_director(IC, param=None):
    """
    Creates a director dict for a PPD.
    
    Parameters
    ----------
    IC : ICobj
    param : .param dict
        If not supplied, param = IC.param
    """    
    if param is None:
        
        param = IC.snapshot_param
        
    # largest radius to plot
    r_director = float(0.9 * IC.pos.r.max())
    # Maximum surface density
    sigma_min = float(IC.sigma(r_director))
    # surface density at largest radius
    sigma_max = float(IC.sigma.input_dict['sigma'].max())
    # Create director dict
    director = make_director(sigma_min, sigma_max, r_director, \
    filename=param['achOutName'])
    
    return director
    
def sink_radius(IC):
    """
    Determine a reasonable sink radius for the star particles depending on 
    the star system type (e.g., single star, binary, etc...)
    
    Parameters
    ----------
    IC : IC object
    
    Returns
    -------
    r_sink : SimArray
        Sink radius for star particles
    """
    
    # Set up the sink radius
    starMode = IC.settings.physical.starMode.lower()
    
    if starMode is 'binary':
        
        binsys = IC.settings.physical.binsys
        #Set Sink Radius to be mass-weighted average of Roche lobes of two stars
        r1 = AddBinary.calcRocheLobe(binsys.m1/binsys.m2,binsys.a) 
        r2 = AddBinary.calcRocheLobe(binsys.m2/binsys.m1,binsys.a)
        p = strip_units(binsys.m1/(binsys.m1 + binsys.m2))
        r_sink = (r1*p) + (r2*(1.0-p))
    
    else:
    
        r_sink = IC.pos.r.min()
        
    return r_sink
    
def setup_sinks(IC, snapshot, param):
    """
    Sets up snapshot and param for stars that are sinks.
    
    Parameters
    ----------
    IC : IC obj
    snapshot : SimSnap
    param : param dict
    
    Returns
    -------
    None
    """
    units = diskpy.pychanga.units_from_param(param)
    # Set the star tforms to a negative number.  This allows UW ChaNGa treat 
    # stars as sink particles
    snapshot.star['tform'] = -1.0   
    # Set sink radius for stars
    r_sink = sink_radius(IC)
    r_sink = strip_units(r_sink)
    param['dSinkBoundOrbitRadius'] = r_sink
    param['dSinkRadius'] = r_sink
    # Set sink mass to be 90% of the smallest star
    Mstar = snapshot.s['mass'].min()
    Mstar.convert_units(units['m_unit'])
    param['dSinkMassMin'] = 0.9 * strip_units(Mstar)
    # Turn sinks on
    param['bDoSinks'] = 1

def init_param(IC, snapshot=None):
    """
    Initializes a ChaNGa param dict (see also diskpy.pychanga.make_param) for
    an IC object.
    
    Parameters
    ----------
    IC : ICobject
        Initial conditions object containing required settings for generating
        the param dict
    snapshot : SimSnap
{'fraction_non_alphanumeric': 0.11609039984547034, 'fraction_numerical': 0.008112806644774966, 'mean_word_length': 2.3446382428940566, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7293790', 'n_tokens_mistral': 3198, 'n_tokens_neox': 3129, 'n_words': 899}