All Downloads are FREE. Search and download functionalities are using the official Maven repository.

au.gov.amsa.util.netcdf.NetCdfWriter Maven / Gradle / Ivy

There is a newer version: 0.6.19
Show newest version
package au.gov.amsa.util.netcdf;

import java.io.File;
import java.io.IOException;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

import com.github.davidmoten.guavamini.Lists;
import com.github.davidmoten.guavamini.Preconditions;

import ucar.ma2.Array;
import ucar.ma2.DataType;
import ucar.ma2.InvalidRangeException;
import ucar.nc2.Attribute;
import ucar.nc2.Dimension;
import ucar.nc2.NetcdfFileWriter;
import ucar.nc2.NetcdfFileWriter.Version;
import ucar.nc2.Variable;

public class NetCdfWriter implements AutoCloseable {

    private final NetcdfFileWriter f;
    private final Map, List> map = new HashMap<>();

    public NetCdfWriter(File file, String version) {
        try {
            f = NetcdfFileWriter.createNew(Version.netcdf3, file.getPath());
            // add version attribute
            f.addGroupAttribute(null, new Attribute("version", version));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    public NetcdfFileWriter writer() {
        return f;
    }

    public NetCdfWriter addAttribute(String name, String value) {
        f.addGroupAttribute(null, new Attribute(name, value));
        return this;
    }

    public  VarBuilder addVariable(String shortName, Class cls) {
        return new VarBuilder(this, shortName, cls);
    }

    public  Var addVariable(String shortName, Optional longName,
            Optional units, Optional encoding, Class cls, int numRecords) {
        Preconditions.checkNotNull(shortName);
        Preconditions.checkNotNull(longName);
        Preconditions.checkNotNull(units);
        Preconditions.checkNotNull(encoding);
        Preconditions.checkNotNull(cls);
        Dimension dimension = f.addDimension(null, shortName, numRecords);
        Variable variable = f.addVariable(null, shortName, toDataType(cls),
                Arrays.asList(dimension));
        if (longName.isPresent())
            variable.addAttribute(new Attribute("long_name", longName.get()));
        if (units.isPresent())
            variable.addAttribute(new Attribute("units", units.get()));
        if (encoding.isPresent())
            variable.addAttribute(new Attribute("encoding", encoding.get()));
        return new Var(this, variable, cls);
    }

    public  NetCdfWriter add(Var variable, T value) {
        @SuppressWarnings("unchecked")
        List list = (List) map.get(variable);
        if (list == null) {
            list = Lists.newArrayList();
            map.put(variable, list);
        }
        list.add(value);
        return this;
    }

    @Override
    public void close() {

        try {
            f.create();
            for (Var var : map.keySet()) {
                List list = map.get(var);
                int[] shape = new int[] { list.size() };
                Array data = Array.factory(DataType.getType(var.cls()), shape);
                for (int i = 0; i < list.size(); i++) {
                    data.setObject(i, list.get(i));
                }
                f.write(var.variable(), data);
            }
            f.close();
        } catch (IOException e) {
            throw new RuntimeException(e);
        } catch (InvalidRangeException e) {
            throw new RuntimeException(e);
        }
    }

    private static DataType toDataType(Class cls) {
        return DataType.DOUBLE;
    }

    public static class VarBuilder {
        private final NetCdfWriter writer;
        final String shortName;
        final Class cls;
        Optional longName = Optional.empty();
        Optional units = Optional.empty();
        Optional encoding = Optional.empty();
        Optional numRecords = Optional.empty();

        VarBuilder(NetCdfWriter writer, String shortName, Class cls) {
            this.writer = writer;
            this.shortName = shortName;
            this.cls = cls;
        }

        public VarBuilder longName(String s) {
            longName = Optional.of(s);
            return this;
        }

        public VarBuilder units(String s) {
            longName = Optional.of(s);
            return this;
        }

        public VarBuilder encoding(String s) {
            longName = Optional.of(s);
            return this;
        }

        public VarBuilder numRecords(int n) {
            this.numRecords = Optional.of(n);
            return this;
        }

        public Var build() {
            return writer.addVariable(shortName, longName, units, encoding, cls, numRecords.get());
        }

    }

    public static class Var {

        private final Variable variable;
        private final Class cls;
        private final NetCdfWriter writer;

        public Var(NetCdfWriter writer, Variable variable, Class cls) {
            this.writer = writer;
            this.variable = variable;
            this.cls = cls;
        }

        public Variable variable() {
            return variable;
        }

        public Class cls() {
            return cls;
        }

        public NetCdfWriter writer() {
            return writer;
        }

        public Var add(T t) {
            writer.add(this, t);
            return this;
        }

    }

}