def plot(arg, *args, **kwargs):
"""
Simple plot function for 1D and 2D coordinate arrays. If the data is complex, the absolute square value of the data will be plottted.
Parameters
-----------
arg: coordinate array
the input data
**kwargs: additional parameters to be passed to the plot functions
Returns
--------
plot: output of ax.plot for 1D and ax.imshow for 2D arrays
"""
import expresso.pycas
import numpy as np
from coordinate_ndarray import CoordinateNDArray
if isinstance(arg, expresso.pycas.Expression):
from .settings import Settings
if len(args) > 0 and isinstance(args[0], Settings):
settings = args[0]
args = list(args)
del args[0]
else:
settings = kwargs.get('settings')
if not settings:
raise ValueError(
'cannot plot expression: no settings provided')
del kwargs['settings']
arg = expression_to_array(arg, settings)
if isinstance(arg, (float, complex)):
print arg
return
elif not isinstance(arg, CoordinateNDArray):
if isinstance(arg, np.ndarray):
pc = expresso.pycas
arg = CoordinateNDArray(
arg, [(pc.S(0), pc.S(n)) for n in arg.shape],
[pc.Symbol('x_%i' % i) for i in range(len(arg.shape))])
else:
raise ValueError(
'cannot plot non CoordinateNDArray object. For plotting regular arrays please use the matplotlib.pyplot module.'
)
if not np.can_cast(arg.data.dtype, np.float128):
if np.all(arg.data.imag == 0): arg = arg.real
else: arg = abs(arg)**2
if len(arg.axis) == 1: return line_plot(arg, *args, **kwargs)
elif len(arg.axis) == 2: return image_plot(arg, *args, **kwargs)
else: raise ValueError("input array must be one or two dimensional")
def evaluator(m):
try:
res = f(m[s.x].value, m[s.y].value)
except:
return False
if res is None:
return False
else:
m[s.z] = pc.S(res)
def __init__(self, settings):
super(FiniteDifferences3D, self).__init__(settings)
from _pypropagate import finite_difference_ACF, finite_difference_A0F
pde = settings.partial_differential_equation
sb = settings.simulation_box
self.__C_is_zero = settings.get_numeric(pde.C) == pc.S(0)
sf = 0.5 if not self.__C_is_zero else 1
z, dz = sb.coordinates[2].symbol, sb.coordinates[2].step
evaluators = self._get_evaluators(
[(pde.ra * sf), (pde.ra * sf).subs(z, z - dz * sf), (pde.rc * sf),
(pde.rc * sf).subs(z, z - dz * sf), (pde.rf * sf),
(pde.rf * sf).subs(z, z - dz * sf), pde.u_boundary,
pde.u_boundary.subs(z, z - dz * sf)],
settings,
return_type=pc.Types.Complex,
compile_to_c=True,
parallel=True)
self.__ra = evaluators[0:2]
self.__rc = evaluators[2:4]
self.__rf = evaluators[4:6]
self.__u_boundary = evaluators[6:8]
self._solver = finite_difference_ACF(
) if not self.__C_is_zero else finite_difference_A0F()
self._solver.resize(self._nx, self._ny)
d, u, l, r = [(self._get_x_indices(), np.zeros(self._nx,
dtype=np.uint)),
(self._get_x_indices(),
np.ones(self._nx, dtype=np.uint) * (self._ny - 1)),
(np.zeros(self._ny,
dtype=np.uint), self._get_y_indices()),
(np.ones(self._ny, dtype=np.uint) * (self._nx - 1),
self._get_y_indices())]
self.__boundary_values = [
np.concatenate([v[0] for v in d, u, l, r]),
np.concatenate([v[1] for v in d, u, l, r]),
np.zeros(2 * self._nx + 2 * self._ny, dtype=np.uint)
]
self._set_initial_field(settings)
self._reset()
def _evaluate(self, expr, settings):
import expresso.pycas as pc
expr = pc.S(expr)
if self.ndim == 1:
sb = settings.simulation_box
pde = settings.partial_differential_equation
x, y, z = pde.coordinates
y0 = getattr(pde, y.name + '0')
expr = settings.get_optimized(expr.subs(y.symbol, y0))
try:
y0i = settings.get_as(y.step.subs(y.symbol, y0), int)
expr = settings.get_optimized(expr.subs(y.index, y0i))
except:
pass
return expr
else:
return settings.get_optimized(expr)
def __init__(self, settings, thread_count=None):
super(Fresnel3D, self).__init__(settings)
if thread_count == None:
import multiprocessing
thread_count = multiprocessing.cpu_count()
self._thread_count = thread_count
pe = settings.partial_differential_equation
x, y, z = settings.partial_differential_equation.coordinates
import expresso.pycas as pc
R, = self._get_evaluators([pc.exp(pe.F * z.step)],
settings,
return_type=pc.Types.Complex,
parallel=not self._F_is_constant_in_z)
if self._F_is_constant_in_z:
self.__R_step = R(*self._get_indices())
else:
self.__R = R
D = pc.numpyfy(
self._evaluate(
pc.exp(
-z.step *
(pe.A * pc.Symbol('kx')**2 + pe.C * pc.Symbol('ky')**2)),
settings))
import numpy as np
fx = 2 * np.pi / (self._nx * self._get_as(x.step, float, settings))
fy = 2 * np.pi / (self._ny * self._get_as(y.step, float, settings))
ky, kx = np.meshgrid(
fy * (self._ny / 2. - np.abs(np.arange(self._ny) - self._ny / 2.)),
fx * (self._nx / 2. - np.abs(np.arange(self._nx) - self._nx / 2.)))
self.__C_is_zero = settings.get_numeric(pe.C) == pc.S(0)
self.__D_step = D(kx=kx, ky=ky, **self._get_indices_dict())
self._set_initial_field(settings)
canonical_form.add_rule(pc.multiplication(s.x), s.x)
canonical_form.add_rule(pc.addition(s.x), s.x)
canonical_form.add_rule(pc.exponentiation(s.x), s.x)
canonical_form.add_rule(1 / s.x, s.x**-1)
canonical_form.add_rule(pc.exp(s.x), pc.e**s.x)
canonical_form.add_rule(s.x**s.y,
s.z,
normalize_exponentiation,
condition=pc.equal(
pc.DominantType(pc.Type(s.y), pc.Types.Real),
pc.Types.Real))
canonical_form.add_rule(pc.exp(s.x), pc.e**s.x)
canonical_form.add_rule(pc.sqrt(s.x), s.x**(1 / pc.S(2)))
canonical_form.add_rule(s.x > s.y, s.y < s.x)
canonical_form.add_rule(s.x >= s.y, s.y <= s.x)
canonical_form.add_rule(s.x <= s.y, pc.Or(s.x < s.y, pc.equal(s.x, s.y)))
canonical_form.add_rule(pc.unequal(s.x, s.y), pc.Not(pc.equal(s.x, s.y)))
canonical_form.add_rule(abs(s.x),
pc.Max(s.x, -s.x),
condition=pc.equal(
pc.DominantType(pc.Type(s.x), pc.Types.Real),
pc.Types.Real))
canonical_form.add_rule(pc.Max(s.a, s.b), -pc.Min(-s.a, -s.b))
canonical_form.add_rule(pc.tan(s.x), pc.sin(s.x) / pc.cos(s.x))
evaluator.add_rule(s.a**s.x * s.b**s.x, (s.a * s.b)**(s.x),
condition=pc.And(
pc.Not(
pc.Or(is_explicit_natural(s.a),
is_explicit_natural(s.b))),
pc.Or(issubtype(s.x, pc.Types.Natural), s.a > 0,
s.b > 0)))
evaluator.add_rule(-(s.x + s.y), -s.x - s.y)
evaluator.add_rule(s.x * -1, -s.x)
evaluator.add_rule(-(-s.x), s.x)
evaluator.add_rule((-s.x) * s.y, -(s.x * s.y))
evaluator.add_rule(1 / -s.x, -(1 / s.x))
evaluator.add_rule(-pc.S(0), 0)
def extract_intersection(m):
ma = pc.MulplicityList(m[s.x], pc.multiplication_group, pc.exponentiation,
pc.real_field)
mb = pc.MulplicityList(m[s.y], pc.multiplication_group, pc.exponentiation,
pc.real_field)
common = ma.intersection(mb)
if len(common) == 0:
return False
m[s.a] = (ma - common).as_expression()
m[s.b] = (mb - common).as_expression()
evaluator.add_rule(pc.OperationType(abs(s.x)), pc.Types.Real)
#evaluator.add_rule(pc.Type(pc.conjugate(s.x)),pc.Type(s.x))
evaluator.add_rule(pc.Type(pc.Indicator(s.x)), pc.Types.Natural)
evaluator.add_rule(pc.Type(pc.OuterPiecewise(s.x)), pc.Type(s.x))
evaluator.add_rule(
pc.Type(pc.InnerPiecewise((s.a, s.b), s.x)),
pc.DominantType(pc.Type(s.a), pc.Type(pc.InnerPiecewise(s.x))))
evaluator.add_rule(pc.Type(pc.InnerPiecewise((s.a, s.b))), pc.Type(s.a))
evaluator.add_rule(pc.Type(pc.derivative(s.x, s.y)), pc.Type(s.x))
evaluator.add_rule(pc.Type(pc.evaluated_at(s.x, s.y, s.z)),
pc.DominantType(pc.Type(s.x), pc.Type(s.z)))
evaluator.add_rule(pc.Type(pc.tmp(s.x)), pc.Type(s.x))
evaluator.add_rule(pc.Type(pc.sqrt(s.x)), pc.Type(s.x**(1 / pc.S(2))))
evaluator.add_rule(
pc.Type(pc.atan2(s.x, s.y)),
pc.DominantType(pc.Type(s.x), pc.Type(s.x), pc.Types.Rational))
for f in [
pc.exp, pc.log, pc.sin, pc.cos, pc.asin, pc.acos, pc.tan, pc.atan,
pc.cot, pc.acot, pc.sinh, pc.cosh, pc.asinh, pc.acosh, pc.tanh,
pc.atanh, pc.coth, pc.acoth
]:
evaluator.add_rule(pc.Type(f(s.x)),
pc.DominantType(pc.Type(s.x), pc.Types.Rational))
def issubtype(x, t):
V = kg * m**2 * s**-3 * A**-1
add_metric_prefixes('V')
J = C * V
add_metric_prefixes('J')
W = J / s
add_metric_prefixes('W')
F = kg**-1 * m**-2 * s**4 * A**2
add_metric_prefixes('F')
ohm = kg * m**2 * s**-3 * A**-2
add_metric_prefixes('ohm')
eV = 1.6021766208 * pc.S(10)**-19 * J
add_metric_prefixes('eV')
h = 6.626070040 * pc.S(10)**-34 * J * s
hbar = h / (2 * pc.pi)
c = 299792458 * m / s
degrees = pc.pi / 180
def contains_unit(expr):
for e in pc.postorder_traversal(expr):
if e in base_units:
return True
return False