001/*-
002 * Copyright 2015, 2016 Diamond Light Source Ltd.
003 *
004 * All rights reserved. This program and the accompanying materials
005 * are made available under the terms of the Eclipse Public License v1.0
006 * which accompanies this distribution, and is available at
007 * http://www.eclipse.org/legal/epl-v10.html
008 */
009
010package org.eclipse.january.dataset;
011
012import org.eclipse.january.DatasetException;
013
014/**
015 * A running mean class
016 */
017public class RunningAverage {
018        
019        private DoubleDataset average;
020        private DoubleDataset sqAveError;
021        private int count = 1;
022
023        /**
024         * @param dataset
025         */
026        public RunningAverage(IDataset dataset) {
027                average = (DoubleDataset) (dataset.getElementClass().equals(Double.class) ? DatasetUtils.convertToDataset(dataset).clone()
028                                : DatasetUtils.cast(DoubleDataset.class, dataset));
029
030                sqAveError = null;
031                Dataset eb = average.getErrorBuffer();
032                if (eb != null) {
033                        sqAveError = DatasetUtils.cast(DoubleDataset.class, eb);
034                }
035        }
036
037        /**
038         * Update average
039         * @param dataset
040         */
041        public void update(IDataset dataset) {
042                count++;
043                IndexIterator it = average.getIterator(true);
044                int[] pos = it.getPos();
045                double f = 1. / count;
046                if (sqAveError == null) {
047                        while (it.hasNext()) {
048                                double m = average.getAbs(it.index);
049                                double v = f * (dataset.getDouble(pos) - m);
050                                average.setAbs(it.index, m + v);
051                        }
052                } else {
053                        double fs = f * f;
054                        double gs = 2 * count - 1;
055                        if (dataset instanceof Dataset) {
056                                final Dataset d = (Dataset) dataset;
057                                final Dataset e = d.getErrorBuffer();
058                                while (it.hasNext()) {
059                                        double m = average.getAbs(it.index);
060                                        double v = f * (d.getDouble(pos) - m);
061                                        average.setAbs(it.index, m + v);
062
063                                        if (e != null) {
064                                                m = sqAveError.getDouble(pos);
065                                                v = fs * (e.getDouble(pos) - gs * m);
066                                                sqAveError.setItem(m + v, pos);
067                                        }
068                                }
069                        } else { // only linear error available
070                                ILazyDataset le = dataset.getErrors();
071                                IDataset e = null;
072                                if (le instanceof IDataset) {
073                                        e = (IDataset) le;
074                                } else if (le != null) {
075                                        try {
076                                                e = le.getSlice();
077                                        } catch (DatasetException e1) {
078                                        }
079                                }
080                                while (it.hasNext()) {
081                                        double m = average.getAbs(it.index);
082                                        double v = f * (dataset.getDouble(pos) - m);
083                                        average.setAbs(it.index, m + v);
084
085                                        if (e != null) {
086                                                m = sqAveError.getDouble(pos);
087                                                v = e.getDouble(pos);
088                                                v = fs * (v * v - gs * m);
089                                                sqAveError.setItem(m + v, pos);
090                                        }
091                                }
092                        }
093                }
094        }
095
096        /**
097         * @return count
098         */
099        public int getCount() {
100                return count;
101        }
102
103        /**
104         * @return current average
105         */
106        public Dataset getCurrentAverage() {
107                if (sqAveError != null) {
108                        Dataset e = sqAveError.clone();
109                        DatasetUtils.makeFinite(e);
110                        average.setErrorBuffer(e);
111                }
112
113                return average;
114        }
115}