def vector(value=[], unit=None):
if unit is None:
if isinstance(value, core.unit):
return VectorQuantity([], unit=value)
elif isinstance(value, ScalarQuantity):
return value.as_vector_with_length(1)
else:
result = AdaptingVectorQuantity()
result.extend(value)
return result
else:
if isinstance(value, ScalarQuantity):
return value.as_vector_with_length(1)
else:
return VectorQuantity(value, unit)
def equal_length_array_or_scalar(array, length=1, mode="continue"):
"""
Returns 'array' if its length is equal to 'length'. If this is not the
case, returns an array of length 'length' with values equal to the first
value of the array (or if 'array' is a scalar, that value. If mode is
"warn", issues a warning if this happens; if mode is "exception" raises an
exception in this case.
"""
try:
array_length = len(array)
if array_length == length:
return array
else:
if mode == "warn":
warnings.warn("Length of array is not equal to %i. Using only\
the first value." % length)
try:
unit = array.unit
value = array[0].value_in(unit)
except:
unit = units.none
value = array[0]
array = VectorQuantity(
array=numpy.ones(length) * value,
unit=unit,
)
return array
elif mode == "exception":
raise Exception("Length of array is not equal to %i. This is\
not supported." % length)
except:
try:
unit = array.unit
value = array.value_in(unit)
except:
unit = units.none
value = array
array = VectorQuantity(
array=numpy.ones(length) * value,
unit=unit,
)
if mode == "warn":
warnings.warn("Using single value for all cases.")
return array
def get_orbital_elements_from_arrays(rel_position_raw,
rel_velocity_raw,
total_masses,
G=nbody_system.G):
"""
Orbital elements from array of relative positions and velocities vectors,
based on orbital_elements_from_binary and adapted to work for arrays (each
line characterises a two body problem).
For circular orbits (eccentricity=0): returns argument of pericenter = 0.,
true anomaly = 0.
For equatorial orbits (inclination=0): longitude of ascending node = 0,
argument of pericenter = arctan2(e_y,e_x).
:argument rel_position: array of vectors of relative positions of the
two-body systems
:argument rel_velocity: array of vectors of relative velocities of the
two-body systems
:argument total_masses: array of total masses for two-body systems
:argument G: gravitational constant
:output semimajor_axis: array of semi-major axes
:output eccentricity: array of eccentricities
:output period: array of orbital periods
:output inc: array of inclinations [radians]
:output long_asc_node: array of longitude of ascending nodes [radians]
:output arg_per_mat: array of argument of pericenters [radians]
:output true_anomaly: array of true anomalies [radians]
"""
if len(numpy.shape(rel_position_raw)) == 1:
rel_position = numpy.zeros([1, 3]) * rel_position_raw[0]
rel_position[0, 0] = rel_position_raw[0]
rel_position[0, 1] = rel_position_raw[1]
rel_position[0, 2] = rel_position_raw[2]
rel_velocity = numpy.zeros([1, 3]) * rel_velocity_raw[0]
rel_velocity[0, 0] = rel_velocity_raw[0]
rel_velocity[0, 1] = rel_velocity_raw[1]
rel_velocity[0, 2] = rel_velocity_raw[2]
else:
rel_position = rel_position_raw
rel_velocity = rel_velocity_raw
separation = (rel_position**2).sum(axis=1)**0.5
n_vec = len(rel_position)
speed_squared = (rel_velocity**2).sum(axis=1)
semimajor_axis = (G * total_masses * separation /
(2. * G * total_masses - separation * speed_squared))
neg_ecc_arg = (
(to_quantity(rel_position).cross(rel_velocity)**2).sum(axis=-1) /
(G * total_masses * semimajor_axis))
filter_ecc0 = (1. <= neg_ecc_arg)
eccentricity = numpy.zeros(separation.shape)
eccentricity[~filter_ecc0] = numpy.sqrt(1.0 - neg_ecc_arg[~filter_ecc0])
eccentricity[filter_ecc0] = 0.
# angular momentum
mom = to_quantity(rel_position).cross(rel_velocity)
# inclination
inc = arccos(mom[:, 2] / to_quantity(mom).lengths())
# Longitude of ascending nodes, with reference direction along x-axis
asc_node_matrix_unit = numpy.zeros(rel_position.shape)
z_vectors = numpy.zeros([n_vec, 3])
z_vectors[:, 2] = 1.
z_vectors = z_vectors | units.none
ascending_node_vectors = z_vectors.cross(mom)
filter_non0_incl = (to_quantity(ascending_node_vectors).lengths().number >
0.)
asc_node_matrix_unit[~filter_non0_incl] = numpy.array([1., 0., 0.])
an_vectors_len = to_quantity(
ascending_node_vectors[filter_non0_incl]).lengths()
asc_node_matrix_unit[filter_non0_incl] = normalize_vector(
ascending_node_vectors[filter_non0_incl], an_vectors_len)
long_asc_node = arctan2(asc_node_matrix_unit[:, 1],
asc_node_matrix_unit[:, 0])
# Argument of periapsis using eccentricity a.k.a. Laplace-Runge-Lenz vector
mu = G * total_masses
pos_unit_vecs = normalize_vector(rel_position, separation)
mom_len = to_quantity(mom).lengths()
mom_unit_vecs = normalize_vector(mom, mom_len)
e_vecs = (normalize_vector(to_quantity(rel_velocity).cross(mom), mu) -
pos_unit_vecs)
# Argument of pericenter cannot be determined for e = 0,
# in this case return 0.0 and 1.0 for the cosines
e_vecs_norm = (e_vecs**2).sum(axis=1)**0.5
filter_non0_ecc = (e_vecs_norm > 1.e-15)
arg_per_mat = VectorQuantity(array=numpy.zeros(long_asc_node.shape),
unit=units.rad)
cos_arg_per = numpy.zeros(long_asc_node.shape)
arg_per_mat[~filter_non0_ecc] = 0. | units.rad
cos_arg_per[~filter_non0_ecc] = 1.
e_vecs_unit = numpy.zeros(rel_position.shape)
e_vecs_unit[filter_non0_ecc] = normalize_vector(
e_vecs[filter_non0_ecc], e_vecs_norm[filter_non0_ecc])
cos_arg_per[filter_non0_ecc] = (e_vecs_unit[filter_non0_ecc] *
asc_node_matrix_unit[filter_non0_ecc]).sum(
axis=-1)
e_cross_an = numpy.zeros(e_vecs_unit.shape)
e_cross_an[filter_non0_ecc] = numpy.cross(
e_vecs_unit[filter_non0_ecc], asc_node_matrix_unit[filter_non0_ecc])
e_cross_an_norm = (e_cross_an**2).sum(axis=1)**0.5
filter_non0_e_cross_an = (e_cross_an_norm != 0.)
ss = -numpy.sign((mom_unit_vecs[filter_non0_e_cross_an] *
e_cross_an[filter_non0_e_cross_an]).sum(axis=-1))
# note change in size in sin_arg_per and cos_arg_per; they are not used further
sin_arg_per = ss * e_cross_an_norm[filter_non0_e_cross_an]
cos_arg_per = cos_arg_per[filter_non0_e_cross_an]
arg_per_mat[filter_non0_e_cross_an] = arctan2(sin_arg_per, cos_arg_per)
# in case longitude of ascending node is 0, omega=arctan2(e_y,e_x)
arg_per_mat[~filter_non0_e_cross_an & filter_non0_ecc] = (arctan2(
e_vecs[~filter_non0_e_cross_an & filter_non0_ecc, 1],
e_vecs[~filter_non0_e_cross_an & filter_non0_ecc, 0]))
filter_negative_zmom = (~filter_non0_e_cross_an
& filter_non0_ecc
& (mom[:, 2] < 0. * mom[0, 0]))
arg_per_mat[filter_negative_zmom] = (2. * numpy.pi -
arg_per_mat[filter_negative_zmom])
# true anomaly
cos_true_anomaly = (e_vecs_unit * pos_unit_vecs).sum(axis=-1)
e_cross_pos = numpy.cross(e_vecs_unit, pos_unit_vecs)
ss2 = numpy.sign((mom_unit_vecs * e_cross_pos).sum(axis=-1))
sin_true_anomaly = ss2 * (e_cross_pos**2).sum(axis=1)**0.5
true_anomaly = arctan2(sin_true_anomaly, cos_true_anomaly)
return (semimajor_axis, eccentricity, true_anomaly, inc, long_asc_node,
arg_per_mat)
def __B_lambda__(l, temp):
tmp = VectorQuantity([], 1e+50 * units.m**-1 * units.kg * units.s**-3)
for t in temp:
curr = 2*h*c**2/l**5*1./(e**(h*c/(l*kB*t))-1)
tmp.append(curr)
return tmp
