From 9431be07d127d451f0b81aca81f6fd6e1360eb68 Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Mon, 10 May 2021 14:43:37 +0200 Subject: [PATCH] Apply suggestions from code review Close #41 --- eddymotion/model.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/eddymotion/model.py b/eddymotion/model.py index 9ce83233..a22ce278 100644 --- a/eddymotion/model.py +++ b/eddymotion/model.py @@ -94,15 +94,15 @@ def predict(self, gradient, **kwargs): class AverageDWModel: """A trivial model that returns an average map.""" - __slots__ = ("_data",) + __slots__ = ("_data", "_gtab") def __init__(self, gtab, **kwargs): """Implement object initialization.""" - return # do nothing at initialization time + self._gtab = gtab def fit(self, data, **kwargs): """Calculate the average.""" - self._data = data.mean(-1) + self._data = data[..., self._gtab[..., 3] > 50].mean(-1) def predict(self, gradient, **kwargs): """Return the average map."""