From ded7a3967c9ae8420a13a8f7372f6a12bbf8e0fd Mon Sep 17 00:00:00 2001 From: rowanc1 Date: Sun, 27 Apr 2014 20:51:03 -0700 Subject: [PATCH] Updates to survey fields class. --- SimPEG/Survey.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/SimPEG/Survey.py b/SimPEG/Survey.py index 7954bc21..34b803d5 100644 --- a/SimPEG/Survey.py +++ b/SimPEG/Survey.py @@ -229,7 +229,7 @@ class Fields(object): dtype = self.dtype[name] else: dtype = self.dtype - field = np.empty(self._storageShape(nP), dtype=dtype) + field = np.zeros(self._storageShape(nP), dtype=dtype) self._fields[name] = field @@ -280,17 +280,13 @@ class Fields(object): assert type(value) is dict, 'New fields must be a dictionary, if field is not specified.' newFields = value elif name in self.knownFields: - assert type(value) is np.ndarray, 'Must be set to a numpy array' newFields = {name: value} else: raise Exception('Unknown setter') for name in newFields: field = self._initStore(name) - NEWF = newFields[name] - if field.shape[1] == 1 or NEWF.ndim == 1: - NEWF = Utils.mkvc(NEWF,2) - self._setField(field, NEWF, ind) + self._setField(field, newFields[name], ind) def __getitem__(self, key): ind, name = self._indexAndNameFromKey(key) @@ -302,6 +298,8 @@ class Fields(object): return self._getField(name, ind) def _setField(self, field, val, ind): + if type(val) is np.ndarray and (field.shape[1] == 1 or val.ndim == 1): + val = Utils.mkvc(val,2) field[:,ind] = val def _getField(self, name, ind): @@ -310,6 +308,9 @@ class Fields(object): out = Utils.mkvc(out) return out + def __contains__(self, other): + return self._fields.__contains__(other) + class TimeFields(Fields): """Fancy Field Storage for time domain problems @@ -342,6 +343,8 @@ class TimeFields(Fields): def _setField(self, field, val, ind): txInd, timeInd = ind + if type(val) is np.ndarray and (val.ndim == 1): + val = Utils.mkvc(val,2) field[:,txInd,timeInd] = val def _getField(self, name, ind):