-
Notifications
You must be signed in to change notification settings - Fork 18
/
Copy pathcore.py
840 lines (676 loc) · 30.3 KB
/
core.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
# gridDataFormats --- python modules to read and write gridded data
# Copyright (c) 2009-2014 Oliver Beckstein <orbeckst@gmail.com>
# Released under the GNU Lesser General Public License, version 3 or later.
"""
:mod:`gridData.core` --- Core functionality for storing n-D grids
=================================================================
The :mod:`core` module contains classes and functions that are
independent of the grid data format. In particular this module
contains the :class:`Grid` class that acts as a universal constructor
for specific formats::
g = Grid(**kwargs) # construct
g.export(filename, format) # export to the desired format
Some formats can also be read::
g = Grid() # make an empty Grid
g.load(filename) # populate with data from filename
Classes and functions
---------------------
"""
# Having consistent truedivision in this module is essential so that
# its behavior is fully consistent in Python 2 and Python 3.
from __future__ import absolute_import, division
import six
from six.moves import cPickle, range, zip
import os
import errno
import numpy
# For interpolated grids: need scipy.ndimage but we import it only when needed:
# import scipy
from . import OpenDX
from . import gOpenMol
from . import mrc
def _grid(x):
"""Access the underlying ndarray of a Grid object or return the object itself"""
try:
return x.grid
except AttributeError:
return x
class Grid(object):
"""Class to manage a multidimensional grid object.
The export(format='dx') method always exports a 3D object, the
rest should work for an array of any dimension.
The grid (Grid.grid) can be manipulated as a standard numpy
array.
The attribute Grid.metadata holds a user-defined dictionary that
can be used to annotate the data. It is saved with save().
"""
default_format = 'DX'
def __init__(self, grid=None, edges=None, origin=None, delta=None,
metadata=None, interpolation_spline_order=3,
file_format=None):
"""
Create a Grid object from data.
From a numpy.histogramdd()::
grid,edges = numpy.histogramdd(...)
g = Grid(grid,edges=edges)
From an arbitrary grid::
g = Grid(grid,origin=origin,delta=delta)
From a saved file::
g = Grid(filename)
or ::
g = Grid()
g.load(filename)
:Arguments:
grid
histogram or density, defined on numpy nD array, or filename
edges
list of arrays, the lower and upper bin edges along the axes
(both are output by numpy.histogramdd())
origin
cartesian coordinates of the center of grid[0,0,...,0]
delta
Either n x n array containing the cell lengths in each dimension,
or n x 1 array for rectangular arrays.
metadata
a user defined dictionary of arbitrary values
associated with the density; the class does not touch
metadata[] but stores it with save()
interpolation_spline_order
order of interpolation function for resampling; cubic splines = 3 [3]
file_format
file format; only necessary when ``grid`` is a filename (see :meth:`Grid.load`);
default is ``None`` and the file format is autodetected.
.. versionchanged:: 0.5.0
New *file_format* keyword argument.
.. versionchanged:: 0.7.0
CCP4 files are now read with :class:`gridData.mrc.MRC` and not anymore
with the deprecated/buggy `ccp4.CCP4`
"""
# file formats are guess from extension == lower case key
self._exporters = {
'DX': self._export_dx,
'PKL': self._export_python,
'PICKLE': self._export_python, # compatibility
'PYTHON': self._export_python, # compatibility
}
self._loaders = {
'CCP4': self._load_mrc,
'MRC': self._load_mrc,
'DX': self._load_dx,
'PLT': self._load_plt,
'PKL': self._load_python,
'PICKLE': self._load_python, # compatibility
'PYTHON': self._load_python, # compatibility
}
self.metadata = metadata if metadata is not None else {}
self.__interpolated = None # cache for interpolated grid
self.__interpolation_spline_order = interpolation_spline_order
self.interpolation_cval = None # default to using min(grid)
if grid is not None:
if isinstance(grid, six.string_types):
# can probably safely try to load() it...
filename = grid
else:
try:
# Can we read this as a file?
# Use str(x) to work with py.path.LocalPath and pathlib.Path instances
# even for Python < 3.6
with open(str(grid), 'rb'):
pass
except (OSError, IOError):
# no, this is probably an array-like thingy
filename = None
else:
# yes, let's use it as a file
filename = str(grid)
if filename is not None:
self.load(filename, file_format=file_format)
else:
self._load(grid, edges, metadata, origin, delta)
@property
def interpolation_spline_order(self):
"""Order of the B-spline interpolation of the data.
3 = cubic; 4 & 5 are also supported
Only choose values that are acceptable to
:func:`scipy.ndimage.spline_filter`!
See Also
--------
interpolated
"""
return self.__interpolation_spline_order
@interpolation_spline_order.setter
def interpolation_spline_order(self, x):
"""Setting the ``interpolation_spline_order`` updates :func:`interpolated`
Because we cache the interpolation function, we need to rebuild the
cache whenever the interpolation order changes: this is
handled by :meth:`_update`
"""
self.__interpolation_spline_order = x
self._update()
def resample(self, edges):
"""Resample data to a new grid with edges *edges*.
This method creates a new grid with the data from the current
grid resampled to a regular grid specified by *edges*. The
order of the interpolation is set by
:attr:`Grid.interpolation_spline_order`: change the value
*before* calling :meth:`resample`.
Parameters
----------
edges : tuple of arrays or Grid
edges of the new grid or a :class:`Grid` instance that
provides :attr:`Grid.edges`
Returns
-------
Grid
a new :class:`Grid` with the data interpolated over the
new grid cells
Examples
--------
Providing *edges* (a tuple of three arrays, indicating the
boundaries of each grid cell)::
g = grid.resample(edges)
As a convenience, one can also supply another :class:`Grid` as
the argument for this method ::
g = grid.resample(othergrid)
and the edges are taken from :attr:`Grid.edges`.
"""
try:
edges = edges.edges # can also supply another Grid
except AttributeError:
pass
midpoints = self._midpoints(edges)
coordinates = ndmeshgrid(*midpoints)
# feed a meshgrid to generate all points
newgrid = self.interpolated(*coordinates)
return self.__class__(newgrid, edges)
def resample_factor(self, factor):
"""Resample to a new regular grid.
Parameters
----------
factor : float
The number of grid cells are scaled with `factor` in each
dimension, i.e., ``factor * N_i`` cells along each
dimension i. Must be positive, and cannot result in fewer
than 2 cells along a dimension.
Returns
-------
Grid
See Also
--------
resample
.. versionchanged:: 0.6.0
Previous implementations would not alter the range of the grid edges
being resampled on. As a result, values at the grid edges would creep
steadily inward. The new implementation recalculates the extent of
grid edges for every resampling.
"""
if float(factor) <= 0:
raise ValueError("Factor must be positive")
# Determine current spacing
spacing = (numpy.array(self._max_edges()) - numpy.array(self._min_edges())) / (
-1 + numpy.array(self._len_edges()))
# First guess at the new spacing is inversely related to the
# magnification factor.
newspacing = spacing / float(factor)
smidpoints = numpy.array(self._midpoints())
# We require that the new spacing result in an even subdivision of the
# existing midpoints
newspacing = (smidpoints[:, -1] - smidpoints[:, 0]) / (numpy.maximum(
1, numpy.floor((smidpoints[:, -1] - smidpoints[:, 0]) / newspacing)))
# How many edge points should there be? It is the number of intervals
# between midpoints + 2
edgelength = 2 + \
numpy.round((smidpoints[:, -1] - smidpoints[:, 0]) / newspacing)
edges = [numpy.linspace(start, stop, num=int(N), endpoint=True) for (start, stop, N) in zip(
smidpoints[:, 0] - 0.5 * newspacing, smidpoints[:, -1] + 0.5 * newspacing, edgelength)]
return self.resample(edges)
def _update(self):
"""compute/update all derived data
Can be called without harm and is idem-potent.
Updates these attributes and methods:
:attr:`origin`
the center of the cell with index 0,0,0
:attr:`midpoints`
centre coordinate of each grid cell
:meth:`interpolated`
spline interpolation function that can generated a value for
coordinate
"""
self.delta = numpy.array(list(
map(lambda e: (e[-1] - e[0]) / (len(e) - 1), self.edges)))
self.midpoints = self._midpoints(self.edges)
self.origin = numpy.array(list(map(lambda m: m[0], self.midpoints)))
if self.__interpolated is not None:
# only update if we are using it
self.__interpolated = self._interpolationFunctionFactory()
@property
def interpolated(self):
"""B-spline function over the data grid(x,y,z).
The :func:`interpolated` function allows one to obtain data
values for any values of the coordinates::
interpolated([x1,x2,...],[y1,y2,...],[z1,z2,...]) -> F[x1,y1,z1],F[x2,y2,z2],...
The interpolation order is set in
:attr:`Grid.interpolation_spline_order`.
The interpolated function is computed once and is cached for better
performance. Whenever :attr:`~Grid.interpolation_spline_order` is
modified, :meth:`Grid.interpolated` is recomputed.
The value for unknown data is set in :attr:`Grid.interpolation_cval`
(TODO: also recompute when ``interpolation_cval`` value is changed.)
Example
-------
Example usage for resampling::
XX, YY, ZZ = numpy.mgrid[40:75:0.5, 96:150:0.5, 20:50:0.5]
FF = interpolated(XX, YY, ZZ)
Note
----
Values are interpolated with a spline function. It is possible
that the spline will generate values that would not normally
appear in the data. For example, a density is non-negative but
a cubic spline interpolation can generate negative values,
especially at the boundary between 0 and high values.
Internally, the function uses :func:`scipy.ndimage.map_coordinates`
with ``mode="constant"`` whereby interpolated values outside
the interpolated grid are determined by filling all values beyond
the edge with the same constant value, defined by the
:attr:`interpolation_cval` parameter, which when not set defaults
to the minimum value in the interpolated grid.
.. versionchanged:: 0.6.0
Interpolation outside the grid is now performed with
``mode="constant"`rather than ``mode="nearest"``, eliminating
extruded volumes when interpolating beyond the grid.
"""
if self.__interpolated is None:
self.__interpolated = self._interpolationFunctionFactory()
return self.__interpolated
def _map_edges(self, func, edges=None):
if edges is None:
edges = self.edges
return [func(e) for e in edges]
def _midpoints(self, edges=None):
return self._map_edges(lambda e: 0.5 * (e[:-1] + e[1:]), edges=edges)
def _len_edges(self, edges=None):
return self._map_edges(len, edges=edges)
def _min_edges(self, edges=None):
return self._map_edges(numpy.min, edges=edges)
def _max_edges(self, edges=None):
return self._map_edges(numpy.max, edges=edges)
def _guess_format(self, filename, file_format=None, export=True):
if export:
available = self._exporters
else:
available = self._loaders
if file_format is None:
splitted = os.path.splitext(filename)
if splitted[1][1:] in ('gz', ):
file_format = os.path.splitext(splitted[0])[1][1:]
else:
file_format = splitted[1][1:]
file_format = file_format.upper()
if not file_format:
file_format = self.default_format
if file_format not in available:
raise ValueError(
"File format {} not available, choose one of {}".format(
file_format, available.keys()))
return file_format
def _get_exporter(self, filename, file_format=None):
return self._exporters[self._guess_format(filename,
file_format=file_format,
export=True)]
def _get_loader(self, filename, file_format=None):
return self._loaders[self._guess_format(filename,
file_format=file_format,
export=False)]
def _load(
self,
grid=None,
edges=None,
metadata=None,
origin=None,
delta=None):
if edges is not None:
# set up from histogramdd-type data
self.grid = numpy.asanyarray(grid)
self.edges = edges
self._update()
elif origin is not None and delta is not None:
# setup from generic data
origin = numpy.asanyarray(origin)
delta = numpy.asanyarray(delta)
if len(origin) != grid.ndim:
raise TypeError(
"Dimension of origin is not the same as grid dimension.")
if delta.shape == () and numpy.isreal(delta):
delta = numpy.ones(grid.ndim) * delta
elif delta.ndim > 1:
raise NotImplementedError(
"Non-rectangular grids are not supported.")
elif len(delta) != grid.ndim:
raise TypeError("delta should be scalar or array-like of"
"len(grid.ndim)")
# note that origin is CENTER so edges must be shifted by -0.5*delta
self.edges = [origin[dim] +
(numpy.arange(m + 1) - 0.5) * delta[dim]
for dim, m in enumerate(grid.shape)]
self.grid = numpy.asanyarray(grid)
self._update()
else:
raise ValueError(
"Wrong/missing data to set up Grid. Use Grid() or "
"Grid(grid=<array>, edges=<list>) or "
"Grid(grid=<array>, origin=(x0, y0, z0), delta=(dx, dy, dz)):\n"
"grid={0} edges={1} origin={2} delta={3}".format(
grid, edges, origin, delta))
def load(self, filename, file_format=None):
"""Load saved (pickled or dx) grid and edges from <filename>.pickle
Grid.load(<filename>.pickle)
Grid.load(<filename>.dx)
The load() method calls the class's constructor method and
completely resets all values, based on the loaded data.
"""
filename = str(filename)
if not os.path.exists(filename):
# check before we try to detect the file type because
# _guess_fileformat() does not work well with things that
# are not really a file
raise IOError(errno.ENOENT, "file not found", filename)
loader = self._get_loader(filename, file_format=file_format)
loader(filename)
def _load_python(self, filename):
with open(filename, 'rb') as f:
saved = cPickle.load(f)
self._load(grid=saved['grid'],
edges=saved['edges'],
metadata=saved['metadata'])
def _load_mrc(self, filename):
"""Initializes Grid from a MRC/CCP4 file."""
mrcfile = mrc.MRC(filename)
grid, edges = mrcfile.histogramdd()
self._load(grid=grid, edges=edges, metadata=self.metadata)
# Store header for access from Grid object (undocumented)
# https://github.com/MDAnalysis/GridDataFormats/pull/100#discussion_r782604833
self._mrc_header = mrcfile.header.copy()
def _load_dx(self, filename):
"""Initializes Grid from a OpenDX file."""
dx = OpenDX.field(0)
dx.read(filename)
grid, edges = dx.histogramdd()
self._load(grid=grid, edges=edges, metadata=self.metadata)
def _load_plt(self, filename):
"""Initialize Grid from gOpenMol plt file."""
g = gOpenMol.Plt()
g.read(filename)
grid, edges = g.histogramdd()
self._load(grid=grid, edges=edges, metadata=self.metadata)
def export(self, filename, file_format=None, type=None, typequote='"'):
"""export density to file using the given format.
The format can also be deduced from the suffix of the filename
though the *format* keyword takes precedence.
The default format for export() is 'dx'. Use 'dx' for
visualization.
Implemented formats:
dx
:mod:`OpenDX`
pickle
pickle (use :meth:`Grid.load` to restore); :meth:`Grid.save`
is simpler than ``export(format='python')``.
Parameters
----------
filename : str
name of the output file
file_format : {'dx', 'pickle', None} (optional)
output file format, the default is "dx"
type : str (optional)
for DX, set the output DX array type, e.g., "double" or "float".
By default (``None``), the DX type is determined from the numpy
dtype of the array of the grid (and this will typically result in
"double").
.. versionadded:: 0.4.0
typequote : str (optional)
For DX, set the character used to quote the type string;
by default this is a double-quote character, '"'.
Custom parsers like the one from NAMD-GridForces (backend for MDFF)
expect no quotes, and typequote='' may be used to appease them.
.. versionadded:: 0.5.0
"""
filename = str(filename)
exporter = self._get_exporter(filename, file_format=file_format)
exporter(filename, type=type, typequote=typequote)
# note: the _export_FORMAT() methods all take the filename as a mandatory
# argument. They can process kwargs but they are not required to do
# so. However, they must ignore any kwargs that they are not processing.
def _export_python(self, filename, **kwargs):
"""Pickle the Grid object
The object is dumped as a dictionary with grid and edges: This
is sufficient to recreate the grid object with __init__().
"""
data = dict(grid=self.grid, edges=self.edges, metadata=self.metadata)
with open(filename, 'wb') as f:
cPickle.dump(data, f, cPickle.HIGHEST_PROTOCOL)
def _export_dx(self, filename, type=None, typequote='"', **kwargs):
"""Export the density grid to an OpenDX file.
The file format is the simplest regular grid array and it is
also understood by VMD's and Chimera's DX reader; PyMOL
requires the dx `type` to be set to "double".
For the file format see
http://opendx.sdsc.edu/docs/html/pages/usrgu068.htm#HDREDF
"""
root, ext = os.path.splitext(filename)
filename = root + '.dx'
comments = [
'OpenDX density file written by gridDataFormats.Grid.export()',
'File format: http://opendx.sdsc.edu/docs/html/pages/usrgu068.htm#HDREDF',
'Data are embedded in the header and tied to the grid positions.',
'Data is written in C array order: In grid[x,y,z] the axis z is fastest',
'varying, then y, then finally x, i.e. z is the innermost loop.']
# write metadata in comments section
if self.metadata:
comments.append('Meta data stored with the python Grid object:')
for k in self.metadata:
comments.append(' ' + str(k) + ' = ' + str(self.metadata[k]))
comments.append(
'(Note: the VMD dx-reader chokes on comments below this line)')
components = dict(
positions=OpenDX.gridpositions(1, self.grid.shape, self.origin,
self.delta),
connections=OpenDX.gridconnections(2, self.grid.shape),
data=OpenDX.array(3, self.grid, type=type, typequote=typequote),
)
dx = OpenDX.field('density', components=components, comments=comments)
if ext == '.gz':
filename = root + ext
dx.write(filename)
def save(self, filename):
"""Save a grid object to <filename>.pickle
Internally, this calls
``Grid.export(filename, format="python")``. A grid can be
regenerated from the saved data with ::
g = Grid(filename="grid.pickle")
.. note::
The pickle format depends on the Python version and
therefore it is not guaranteed that a grid saved with, say,
Python 2.7 can also be read with Python 3.5. The OpenDX format
is a better alternative for portability.
"""
self.export(filename, file_format="pickle")
def centers(self):
"""Returns the coordinates of the centers of all grid cells as an
iterator."""
for idx in numpy.ndindex(self.grid.shape):
yield self.delta * numpy.array(idx) + self.origin
def check_compatible(self, other):
"""Check if *other* can be used in an arithmetic operation.
1) *other* is a scalar
2) *other* is a grid defined on the same edges
:Raises: :exc:`TypeError` if not compatible.
"""
if not (numpy.isreal(other) or self == other):
raise TypeError(
"The argument can not be arithmetically combined with the grid. "
"It must be a scalar or a grid with identical edges. "
"Use Grid.resample(other.edges) to make a new grid that is "
"compatible with other.")
return True
def _interpolationFunctionFactory(self, spline_order=None, cval=None):
"""Returns a function F(x,y,z) that interpolates any values on the grid.
_interpolationFunctionFactory(self,spline_order=3,cval=None) --> F
*cval* is set to :meth:`Grid.grid.min`. *cval* cannot be chosen too
large or too small or NaN because otherwise the spline interpolation
breaks down near that region and produces wild oscillations.
.. Note:: Only correct for equally spaced values (i.e. regular edges with
constant delta).
.. SeeAlso:: http://www.scipy.org/Cookbook/Interpolation
"""
# for scipy >=0.9: should use scipy.interpolate.griddata
# http://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.griddata.html#scipy.interpolate.griddata
# (does it work for nD?)
import scipy.ndimage
if spline_order is None:
# must be compatible with whatever
# :func:`scipy.ndimage.spline_filter` takes.
spline_order = self.interpolation_spline_order
if cval is None:
cval = self.interpolation_cval
data = self.grid
if cval is None:
cval = data.min()
try:
# masked arrays, fill with min: should keep spline happy
_data = data.filled(cval)
except AttributeError:
_data = data
coeffs = scipy.ndimage.spline_filter(_data, order=spline_order)
x0 = self.origin
dx = self.delta
def _transform(cnew, c0, dc):
return (numpy.atleast_1d(cnew) - c0) / dc
def interpolatedF(*coordinates):
"""B-spline function over the data grid(x,y,z).
interpolatedF([x1,x2,...],[y1,y2,...],[z1,z2,...]) -> F[x1,y1,z1],F[x2,y2,z2],...
Example usage for resampling::
>>> XX,YY,ZZ = numpy.mgrid[40:75:0.5, 96:150:0.5, 20:50:0.5]
>>> FF = _interpolationFunction(XX,YY,ZZ)
"""
_coordinates = numpy.array(
[_transform(coordinates[i], x0[i], dx[i]) for i in range(len(
coordinates))])
return scipy.ndimage.map_coordinates(coeffs,
_coordinates,
prefilter=False,
mode='constant',
cval=cval)
return interpolatedF
def __eq__(self, other):
if not isinstance(other, Grid):
return False
return numpy.all(
other.grid == self.grid) and numpy.all(
other.origin == self.origin) and numpy.all(
numpy.all(
other_edge == self_edge) for other_edge,
self_edge in zip(
other.edges,
self.edges))
def __ne__(self, other):
return not self.__eq__(other)
def __add__(self, other):
self.check_compatible(other)
return self.__class__(self.grid + _grid(other), edges=self.edges)
def __sub__(self, other):
self.check_compatible(other)
return self.__class__(self.grid - _grid(other), edges=self.edges)
def __mul__(self, other):
self.check_compatible(other)
return self.__class__(self.grid * _grid(other), edges=self.edges)
def __truediv__(self, other):
# truediv will always do true division (in Python 2 and Python 3);
# we use from __future__ include division everywhere
self.check_compatible(other)
return self.__class__(self.grid / _grid(other), edges=self.edges)
def __div__(self, other):
# in Python 2 only (without __future__.division): will do "classic division"
# https://docs.python.org/2/reference/datamodel.html#object.__div__
if not six.PY2:
raise NotImplementedError(
"__div__ is only available in Python 2, use __truediv__")
self.check_compatible(other)
return self.__class__(
self.grid.__div__(
_grid(other)),
edges=self.edges)
def __floordiv__(self, other):
self.check_compatible(other)
return self.__class__(self.grid // _grid(other), edges=self.edges)
def __pow__(self, other):
self.check_compatible(other)
return self.__class__(
numpy.power(
self.grid,
_grid(other)),
edges=self.edges)
def __radd__(self, other):
self.check_compatible(other)
return self.__class__(_grid(other) + self.grid, edges=self.edges)
def __rsub__(self, other):
self.check_compatible(other)
return self.__class__(_grid(other) - self.grid, edges=self.edges)
def __rmul__(self, other):
self.check_compatible(other)
return self.__class__(_grid(other) * self.grid, edges=self.edges)
def __rtruediv__(self, other):
self.check_compatible(other)
return self.__class__(_grid(other) / self.grid, edges=self.edges)
def __rdiv__(self, other):
# in Python 2 only (without __future__.division): will do "classic division"
# https://docs.python.org/2/reference/datamodel.html#object.__div__
if not six.PY2:
raise NotImplementedError(
"__rdiv__ is only available in Python 2, use __rtruediv__")
self.check_compatible(other)
return self.__class__(
self.grid.__rdiv__(
_grid(other)),
edges=self.edges)
def __rfloordiv__(self, other):
self.check_compatible(other)
return self.__class__(_grid(other) // self.grid, edges=self.edges)
def __rpow__(self, other):
self.check_compatible(other)
return self.__class__(
numpy.power(
_grid(other),
self.grid),
edges=self.edges)
def __repr__(self):
try:
bins = self.grid.shape
except AttributeError:
bins = "no"
return '<{0} with {1!r} bins>'.format(self.__class__, bins)
def ndmeshgrid(*arrs):
"""Return a mesh grid for N dimensions.
The input are N arrays, each of which contains the values along one axis of
the coordinate system. The arrays do not have to have the same number of
entries. The function returns arrays that can be fed into numpy functions
so that they produce values for *all* points spanned by the axes *arrs*.
Original from
http://stackoverflow.com/questions/1827489/numpy-meshgrid-in-3d and fixed.
.. SeeAlso: :func:`numpy.meshgrid` for the 2D case.
"""
# arrs = tuple(reversed(arrs)) <-- wrong on stackoverflow.com
arrs = tuple(arrs)
lens = list(map(len, arrs))
dim = len(arrs)
sz = 1
for s in lens:
sz *= s
ans = []
for i, arr in enumerate(arrs):
slc = [1] * dim
slc[i] = lens[i]
arr2 = numpy.asanyarray(arr).reshape(slc)
for j, sz in enumerate(lens):
if j != i:
arr2 = arr2.repeat(sz, axis=j)
ans.append(arr2)
return tuple(ans)