def sample(steps, cfield, mgrid, lcar = None, **kwds):
'''
Sample current field cfield along steps at a granularity near lcar.
If lcar not given divine it based on mgrid.
Return an array of 5-arrays: (x,y,z,time,current)
'''
import larf.util
linspaces = mgrid_to_linspace(mgrid)
curfield = Scalar(cfield, linspaces)
if not lcar:
lcar = min(linspaces[0][1] - linspaces[0][0],
linspaces[1][1] - linspaces[1][0],
linspaces[2][1] - linspaces[2][0])
result = list()
for istep, p1 in enumerate(steps[:-1]):
p2 = steps[istep+1]
step = p2-p1
dist = math.sqrt(numpy.dot(step[:3], step[:3]))
n = max(2, int(dist/lcar))
for count in range(n):
p = p1 + float(count)/float(n)*step
r = p[:3]
cur = curfield(r)
pw = numpy.hstack((p, cur))
result.append(pw)
ret = numpy.asarray(result)
return ret
def _save_mgrid_vtk(result, outfile, **kwds):
'''
Save a result defined on a meshgrid to a VTK file
'''
from tvtk.api import tvtk
from tvtk.api import write_data
arrs = result.array_data_by_type()
mgrid = arrs['mgrid']
shape = mgrid[0].shape
linspaces = mgrid_to_linspace(mgrid, expand=False)
origin = list()
spacing = list()
for ls in linspaces:
print "ls: %s" % str(ls)
origin.append(ls[0])
spacing.append((ls[1] - ls[0]) / ls[2])
print 'origin: %s' % str(origin)
print 'spacing: %s' % str(spacing)
print 'dimensions: %s' % str(shape)
dat = tvtk.ImageData(spacing=spacing, origin=origin, dimensions=shape)
scalar = arrs.get('gscalar', None)
if scalar is not None:
dat.point_data.scalars = scalar.ravel(order='F')
dat.point_data.scalars.name = "gscalar" # fixme, should use name, not type?
vector = arrs.get('gvector', None)
if vector is not None:
dat.point_data.vectors = numpy.asarray(
[vector[i].ravel(order='F') for i in range(3)]).T
dat.point_data.vectors.name = "gvector"
write_data(dat, outfile)
return
def _save_mgrid_vtk(result, outfile, **kwds):
'''
Save a result defined on a meshgrid to a VTK file
'''
from tvtk.api import tvtk
from tvtk.api import write_data
arrs = result.array_data_by_type()
mgrid = arrs['mgrid']
shape = mgrid[0].shape
linspaces = mgrid_to_linspace(mgrid, expand=False)
origin = list()
spacing = list()
for ls in linspaces:
print "ls: %s" % str(ls)
origin.append(ls[0])
spacing.append((ls[1]-ls[0])/ls[2])
print 'origin: %s' % str(origin)
print 'spacing: %s' % str(spacing)
print 'dimensions: %s' % str(shape)
dat = tvtk.ImageData(spacing=spacing, origin=origin, dimensions=shape)
scalar = arrs.get('gscalar',None)
if scalar is not None:
dat.point_data.scalars = scalar.ravel(order='F')
dat.point_data.scalars.name = "gscalar" # fixme, should use name, not type?
vector = arrs.get('gvector',None)
if vector is not None:
dat.point_data.vectors = numpy.asarray([vector[i].ravel(order='F') for i in range(3)]).T
dat.point_data.vectors.name = "gvector"
write_data(dat, outfile)
return
def cmd_oldcurrent(ctx, stepping, weight, charge, name):
'''
Produce currents along steps.
'''
from larf.models import Result, Array
from larf.vector import Scalar
from larf.util import mgrid_to_linspace
import larf.current
ses = ctx.obj['session']
sres= larf.store.result_typed(ses, 'stepping', stepping)
wres= larf.store.result_typed(ses, 'raster', weight)
warr = wres.array_data_by_type()
linspaces = mgrid_to_linspace(warr['mgrid'])
weight = Scalar(warr['gscalar'], linspaces)
# fixme: this needs to go in a module with the usual "tocall" pattern.
arrays = list()
for pname, path in sorted(sres.array_data_by_name().items()):
wf = larf.current.stepwise(weight, path, charge)
arrays.append(Array(name=pname, type='pscalar', data=wf))
res = Result(name=name, type='current', parents=[sres, wres],
params=dict(),
arrays=arrays)
ses.add(res)
ses.flush()
ses.commit()
announce_result('current', res)
return
def patch(vfield, mgrid,
corner = (19*units.mm, 10*units.mm, 10*units.mm),
sep = (0.1*units.mm, 0.1*units.mm), start_time=0.0,
namepat = "path%04d", stepper = 'rkck', **kwds):
'''
Return paths stepped through vfield starting at many points on a
rectangular patch. A collection of (name, path) pairs are
returned. Each path is a numpy array of (x,y,z,t).
'''
# fixme: hardcodes.
lcar=0.1*units.mm
stuck=0.01*units.mm
import larf.util
linspaces = mgrid_to_linspace(mgrid)
velo = Field(vfield, linspaces)
def velocity(notused, r):
return velo(r)
step_fun = 'larf.stepping.step_%s' % stepper
step_fun = larf.util.get_method(step_fun)
stepper = Stepper(velocity, lcar=lcar, step_fun=step_fun, **kwds)
cx,cy,cz = corner
sepy, sepz = sep
x = cx
paths = list()
count = 0
for y in numpy.linspace(-cy, cy, 1+int(2.0*cy/sepy)):
for z in numpy.linspace(-cz, cz, 1+int(2.0*cz/sepz)):
name = namepat % count
count += 1
position = (x,y,z)
visitor = stepper(start_time, position, CollectSteps(StuckDetection(distance=stuck)))
paths.append((name, visitor.array))
continue # z
continue # y
return paths
def cmd_velocity(ctx, method, raster, name):
'''
Calculation a velocity vector field.
'''
from larf.models import Result, Array
from larf.util import mgrid_to_linspace
if method == 'drift': # pretend like I actually give an option!
methname = 'larf.drift.velocity'
else:
click.echo('Unknown velocity calculation method: "%s"' % method)
return 1
ses = ctx.obj['session']
import larf.store
potres = larf.store.result_typed(ses, 'raster', raster)
arrs = potres.array_data_by_type()
potential = arrs['gscalar']
mgrid = arrs['mgrid']
linspaces = mgrid_to_linspace(mgrid, expand = False)
import larf.util
meth = larf.util.get_method(methname)
par = ctx.obj['params']
velo = meth(potential, linspaces, **par)
res = Result(name=name, type='velocity', parents=[potres],
params=dict(method=methname, params=par),
arrays = [
Array(name='domain', type='mgrid', data=mgrid),
Array(name='velocity', type='gvector', data = numpy.asarray(velo))])
ses.add(res)
ses.flush()
resid = res.id
ses.commit()
announce_result('velocity', res)
def patch(vfield, cfield, mgrid,
corner = (19*units.mm, 10*units.mm, 10*units.mm), sep = (0.1*units.mm, 0.1*units.mm), start_time=0.0,
digi_time=0.0, tick=0.5*units.ms, nticks=100, **kwds):
'''
Return waveforms sampled from cfield from steps starting at many
points on a rectangular patch.
'''
# fixme: hardcodes.
lcar=0.1*units.mm
stuck=0.01*units.mm
import larf.util
linspaces = mgrid_to_linspace(mgrid)
velo = Field(vfield, linspaces)
def velocity(notused, r):
return velo(r)
from larf.stepping import step_rkck
stepper = Stepper(velocity, lcar=lcar, step_fun=step_rkck)
cx,cy,cz = corner
sepy, sepz = sep
waveforms = list()
step_points = list()
x = cx
for y in numpy.arange(-cy, cy, sepy):
for z in numpy.arange(-cz, cz, sepz):
position = (x,y,z)
#print position
visitor = stepper(start_time, position, CollectSteps(StuckDetection(distance=stuck)))
steps = larf.current.sample(visitor.array, cfield, mgrid, lcar)
waveform = numpy.asarray([0.0]*nticks, dtype=float)
count = numpy.asarray([0]*nticks, dtype=int)
ptt = numpy.asarray([0.0]*nticks, dtype=float)
ptx = numpy.asarray([0.0]*nticks, dtype=float)
pty = numpy.asarray([0.0]*nticks, dtype=float)
ptz = numpy.asarray([0.0]*nticks, dtype=float)
def fill(x,y,z,t,cur):
tbin = int((t-digi_time)/tick)
if tbin<0: # underflow
tbin=0
if tbin>=nticks: # overflow
tbin=nticks-1
count[tbin] += 1
waveform[tbin] += cur
ptt[tbin] += t
ptx[tbin] += x
pty[tbin] += y
ptz[tbin] += z
for step in steps:
fill(*step)
for ind in range(len(count)):
n = count[ind]
if n == 0:
continue
waveform[ind] /= n
ptt[ind] /= n
ptx[ind] /= n
pty[ind] /= n
ptz[ind] /= n
pts = numpy.vstack((ptt, ptx, pty, ptz)).T
step_points.append(pts)
waveforms.append(waveform)
continue # z
continue # y
return (numpy.asarray(step_points), numpy.vstack(waveforms))
def patch(vfield,
cfield,
mgrid,
corner=(19 * units.mm, 10 * units.mm, 10 * units.mm),
sep=(0.1 * units.mm, 0.1 * units.mm),
start_time=0.0,
digi_time=0.0,
tick=0.5 * units.ms,
nticks=100,
**kwds):
'''
Return waveforms sampled from cfield from steps starting at many
points on a rectangular patch.
'''
# fixme: hardcodes.
lcar = 0.1 * units.mm
stuck = 0.01 * units.mm
import larf.util
linspaces = mgrid_to_linspace(mgrid)
velo = Field(vfield, linspaces)
def velocity(notused, r):
return velo(r)
from larf.stepping import step_rkck
stepper = Stepper(velocity, lcar=lcar, step_fun=step_rkck)
cx, cy, cz = corner
sepy, sepz = sep
waveforms = list()
step_points = list()
x = cx
for y in numpy.arange(-cy, cy, sepy):
for z in numpy.arange(-cz, cz, sepz):
position = (x, y, z)
#print position
visitor = stepper(start_time, position,
CollectSteps(StuckDetection(distance=stuck)))
steps = larf.current.sample(visitor.array, cfield, mgrid, lcar)
waveform = numpy.asarray([0.0] * nticks, dtype=float)
count = numpy.asarray([0] * nticks, dtype=int)
ptt = numpy.asarray([0.0] * nticks, dtype=float)
ptx = numpy.asarray([0.0] * nticks, dtype=float)
pty = numpy.asarray([0.0] * nticks, dtype=float)
ptz = numpy.asarray([0.0] * nticks, dtype=float)
def fill(x, y, z, t, cur):
tbin = int((t - digi_time) / tick)
if tbin < 0: # underflow
tbin = 0
if tbin >= nticks: # overflow
tbin = nticks - 1
count[tbin] += 1
waveform[tbin] += cur
ptt[tbin] += t
ptx[tbin] += x
pty[tbin] += y
ptz[tbin] += z
for step in steps:
fill(*step)
for ind in range(len(count)):
n = count[ind]
if n == 0:
continue
waveform[ind] /= n
ptt[ind] /= n
ptx[ind] /= n
pty[ind] /= n
ptz[ind] /= n
pts = numpy.vstack((ptt, ptx, pty, ptz)).T
step_points.append(pts)
waveforms.append(waveform)
continue # z
continue # y
return (numpy.asarray(step_points), numpy.vstack(waveforms))