sr: Add OutputVariables to SRPredict.

Now output variables can be specified using equations just like in 
regular InMAP.

cloud/config_test.go

@@ -124,6 +124,7 @@ func TestRunInputFromViper(t *testing.T) {
 
 func TestSRPredictInputFromViper(t *testing.T) {
 	cfg := inmaputil.InitializeConfig()
+	cfg.Set("OutputVariables", "{\"PrimPM25\":\"PrimaryPM25\"}")
 	js, err := cloud.JobSpec(cfg.Root, cfg.Viper, "test_job", []string{"srpredict"}, cfg.InputFiles(), 1)
 	if err != nil {
 		t.Fatal(err)
@@ -140,6 +141,7 @@ func TestSRPredictInputFromViper(t *testing.T) {
 		"--EmissionUnits":       "tons/year",
 		"--EmissionsShapefiles": "258bbcefe8c0073d6f323351463be9e9685e74bb92e367ca769b9536ed247213.shp",
 		"--OutputFile":          "inmap_output.shp",
+		"--OutputVariables":     "{\"PrimPM25\":\"PrimaryPM25\"}",
 		"--SR.OutputFile":       "${INMAP_ROOT_DIR}/cmd/inmap/testdata/output_${InMAPRunType}.shp",
 		"--VarGrid.GridProj":    "+proj=lcc +lat_1=33.000000 +lat_2=45.000000 +lat_0=40.000000 +lon_0=-97.000000 +x_0=0 +y_0=0 +a=6370997.000000 +b=6370997.000000 +to_meter=1",
 	}

cmd/inmap/testdata/output_SRPredict.prj

@@ -1 +1 @@
-PROJCS["Lambert_Conformal_Conic",GEOGCS["GCS_unnamed ellipse",DATUM["D_unknown",SPHEROID["Unknown",6370997,0]],PRIMEM["Greenwich",0],UNIT["Degree",0.017453292519943295]],PROJECTION["Lambert_Conformal_Conic"],PARAMETER["standard_parallel_1",33],PARAMETER["standard_parallel_2",45],PARAMETER["latitude_of_origin",40],PARAMETER["central_meridian",-97],PARAMETER["false_easting",0],PARAMETER["false_northing",0],UNIT["Meter",1]]
+PROJCS["Lambert_Conformal_Conic",GEOGCS["GCS_unnamed ellipse",DATUM["D_unknown",SPHEROID["Unknown",6370997.000000,0]],PRIMEM["Greenwich",0],UNIT["Degree",0.017453292519943295]],PROJECTION["Lambert_Conformal_Conic"],PARAMETER["standard_parallel_1",33],PARAMETER["standard_parallel_2",45],PARAMETER["latitude_of_origin",40],PARAMETER["central_meridian",-97],PARAMETER["false_easting",0],PARAMETER["false_northing",0],UNIT["Meter",1]]

docs/cmd/inmap_srpredict.md

@@ -15,16 +15,8 @@ predict uses the SR matrix specified in the configuration file
 	from the emissions specified in the EmissionsShapefiles field in the configuration
 	file, outputting the results in the shapefile specified in OutputFile field.
 	of the configuration file. The EmissionUnits field in the configuration
-	file specifies the units of the emissions. Output units are μg particulate
-	matter per m³ air.
-
-		Output variables:
-		PNH4: Particulate ammonium
-		PNO3: Particulate nitrate
-		PSO4: Particulate sulfate
-		SOA: Secondary organic aerosol
-		PrimaryPM25: Primarily emitted PM2.5
-		TotalPM25: The sum of the above components
+	file specifies the units of the emissions. The OutputVariables configuration
+	variable specifies the information to be output.
 
 ```
 inmap srpredict [flags]
@@ -48,6 +40,9 @@ inmap srpredict [flags]
       --OutputFile string             
                                                     OutputFile is the path to the desired output shapefile location. It can
                                                     include environment variables. (default "inmap_output.shp")
+      --OutputVariables string        
+                                                    OutputVariables specifies which model variables should be included in the
+                                                    output file. It can include environment variables. (default "{\"TotalPM25\":\"PrimaryPM25 + pNH4 + pSO4 + pNO3 + SOA\",\"TotalPopD\":\"(exp(log(1.078)/10 * TotalPM25) - 1) * TotalPop * AllCause / 100000\"}\n")
       --SR.OutputFile string          
                                                     SR.OutputFile is the path where the output file is or should be created
                                                      when creating a source-receptor matrix. It can contain environment variables. (default "${INMAP_ROOT_DIR}/cmd/inmap/testdata/output_${InMAPRunType}.shp")

framework.go

@@ -76,9 +76,9 @@ type InMAP struct {
 	// boundary cells; assume bottom boundary is the same as lowest layer
 	topBoundary *cellList
 
-	// popIndices gives the array index of each population type in the PopData
+	// PopIndices gives the array index of each population type in the PopData
 	// field in each Cell.
-	popIndices map[string]int
+	PopIndices map[string]int
 
 	// mortIndices gives the array index of each mortality rate in the mortData
 	// field in each Cell.
@@ -476,7 +476,7 @@ func (d *InMAP) VerticalProfile(variable string, p geom.Point, m Mechanism) (hei
 	}
 	i := 0
 	for !c.boundary {
-		vals[i] = c.getValue(variable, d.popIndices, d.mortIndices, m)
+		vals[i] = c.getValue(variable, d.PopIndices, d.mortIndices, m)
 		height[i] = c.LayerHeight + c.Dz/2.
 		c = (*c.above)[0].Cell
 		i++

inmaputil/cmd.go

@@ -430,16 +430,8 @@ func InitializeConfig() *Cfg {
 	from the emissions specified in the EmissionsShapefiles field in the configuration
 	file, outputting the results in the shapefile specified in OutputFile field.
 	of the configuration file. The EmissionUnits field in the configuration
-	file specifies the units of the emissions. Output units are μg particulate
-	matter per m³ air.
-
-		Output variables:
-		PNH4: Particulate ammonium
-		PNO3: Particulate nitrate
-		PSO4: Particulate sulfate
-		SOA: Secondary organic aerosol
-		PrimaryPM25: Primarily emitted PM2.5
-		TotalPM25: The sum of the above components`,
+	file specifies the units of the emissions. The OutputVariables configuration
+	variable specifies the information to be output.`,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			outChan := outChan()
 
@@ -451,6 +443,10 @@ func InitializeConfig() *Cfg {
 			if err != nil {
 				return err
 			}
+			outputVars, err := checkOutputVars(GetStringMapString("OutputVariables", cfg.Viper))
+			if err != nil {
+				return err
+			}
 			emisUnits, err := checkEmissionUnits(cfg.GetString("EmissionUnits"))
 			if err != nil {
 				return err
@@ -466,6 +462,7 @@ func InitializeConfig() *Cfg {
 				emisUnits,
 				os.ExpandEnv(cfg.GetString("SR.OutputFile")),
 				outputFile,
+				outputVars,
 				shapeFiles,
 				vgc,
 			)
@@ -776,7 +773,7 @@ func InitializeConfig() *Cfg {
 				"TotalPM25": "PrimaryPM25 + pNH4 + pSO4 + pNO3 + SOA",
 				"TotalPopD": "(exp(log(1.078)/10 * TotalPM25) - 1) * TotalPop * AllCause / 100000",
 			},
-			flagsets: []*pflag.FlagSet{cfg.runCmd.PersistentFlags(), cfg.cloudStartCmd.Flags()},
+			flagsets: []*pflag.FlagSet{cfg.runCmd.PersistentFlags(), cfg.cloudStartCmd.Flags(), cfg.srPredictCmd.Flags()},
 		},
 		{
 			name: "NumIterations",

inmaputil/inmap_test.go

@@ -24,6 +24,8 @@ import (
 	"os"
 	"reflect"
 	"testing"
+
+	"github.com/spatialmodel/inmap"
 )
 
 // Set up directory location for configuration files.
@@ -47,7 +49,8 @@ func TestInMAPStaticCreateGrid(t *testing.T) {
 	os.Setenv("InMAPRunType", "static")
 	cfg.Set("config", "../cmd/inmap/configExample.toml")
 	cfg.Root.SetArgs([]string{"run", "steady"})
-	defer os.Remove("../cmd/inmap/testdata/output_static.log")
+	defer os.Remove("inmap_output.log")
+	defer inmap.DeleteShapefile(cfg.GetString("OutputFile"))
 	if err := cfg.Root.Execute(); err != nil {
 		t.Fatal(err)
 	}
@@ -60,7 +63,8 @@ func TestInMAPStaticLoadGrid(t *testing.T) {
 	os.Setenv("InMAPRunType", "staticLoadGrid")
 	cfg.Set("config", "../cmd/inmap/configExample.toml")
 	cfg.Root.SetArgs([]string{"run", "steady"})
-	defer os.Remove("../cmd/inmap/testdata/output_staticLoadGrid.log")
+	defer os.Remove("inmap_output.log")
+	defer inmap.DeleteShapefile(cfg.GetString("OutputFile"))
 	if err := cfg.Root.Execute(); err != nil {
 		t.Fatal(err)
 	}
@@ -73,7 +77,8 @@ func TestInMAPDynamic(t *testing.T) {
 	os.Setenv("InMAPRunType", "dynamic")
 	cfg.Set("config", "../cmd/inmap/configExample.toml")
 	cfg.Root.SetArgs([]string{"run", "steady"})
-	defer os.Remove("../cmd/inmap/testdata/output_dynamic.log")
+	defer os.Remove("inmap_output.log")
+	defer inmap.DeleteShapefile(cfg.GetString("OutputFile"))
 	if err := cfg.Root.Execute(); err != nil {
 		t.Fatal(err)
 	}
@@ -94,7 +99,8 @@ func TestInMAPDynamicRemote_http(t *testing.T) {
 	os.Setenv("InMAPRunType", "dynamic")
 	cfg.Set("config", "../cmd/inmap/configExampleRemote.toml")
 	cfg.Root.SetArgs([]string{"run", "steady"})
-	defer os.Remove("../cmd/inmap/testdata/output_dynamic.log")
+	defer os.Remove("inmap_output.log")
+	defer inmap.DeleteShapefile(cfg.GetString("OutputFile"))
 	if err := cfg.Root.Execute(); err != nil {
 		t.Fatal(err)
 	}
@@ -113,7 +119,8 @@ func TestInMAPDynamicRemote_bucket(t *testing.T) {
 	os.Setenv("InMAPRunType", "dynamic")
 	cfg.Set("config", "../cmd/inmap/configExampleRemote.toml")
 	cfg.Root.SetArgs([]string{"run", "steady"})
-	defer os.Remove("../cmd/inmap/testdata/output_dynamic.log")
+	defer os.Remove("inmap_output.log")
+	defer inmap.DeleteShapefile(cfg.GetString("OutputFile"))
 	if err := cfg.Root.Execute(); err != nil {
 		t.Fatal(err)
 	}

inmaputil/sr.go

@@ -23,10 +23,7 @@ import (
 	"fmt"
 	"log"
 	"os"
-	"path/filepath"
 
-	"github.com/ctessum/geom"
-	"github.com/ctessum/geom/encoding/shp"
 	"github.com/spatialmodel/inmap"
 	"github.com/spatialmodel/inmap/cloud/cloudrpc"
 	"github.com/spatialmodel/inmap/sr"
@@ -112,9 +109,10 @@ func CleanSR(ctx context.Context, jobName, VariableGridData string, VarGrid *inm
 // SRPredict uses the SR matrix specified in SROutputFile
 // to predict concentrations resulting
 // from the emissions in EmissionsShapefiles, outputting the
-// results in OutputFile. EmissionUnits specifies the units
+// results specified by outputVaraibles in OutputFile.
+// EmissionUnits specifies the units
 // of the emissions. VarGrid specifies the variable resolution grid.
-func SRPredict(EmissionUnits, SROutputFile, OutputFile string, EmissionsShapefiles []string, VarGrid *inmap.VarGridConfig) error {
+func SRPredict(EmissionUnits, SROutputFile, OutputFile string, outputVariables map[string]string, EmissionsShapefiles []string, VarGrid *inmap.VarGridConfig) error {
 	msgLog := make(chan string)
 	go func() {
 		for {
@@ -147,52 +145,19 @@ func SRPredict(EmissionUnits, SROutputFile, OutputFile string, EmissionsShapefil
 			return err
 		}
 	}
-	type rec struct {
-		geom.Polygon
-		PNH4, PNO3, PSO4, SOA, PrimaryPM25, TotalPM25 float64
-	}
-
-	var upload uploader
-
-	o, err := shp.NewEncoder(upload.maybeUpload(OutputFile), rec{})
-	if err != nil {
+	if err = r.SetConcentrations(conc); err != nil {
 		return err
 	}
 
+	var upload uploader
+	o := upload.maybeUpload(OutputFile)
 	if upload.err != nil {
 		return upload.err
 	}
 
-	g := r.Geometry()
-
-	for i, tpm := range conc.TotalPM25() {
-		r := rec{
-			Polygon:     g[i].(geom.Polygon),
-			PNH4:        conc.PNH4[i],
-			PNO3:        conc.PNO3[i],
-			PSO4:        conc.PSO4[i],
-			PrimaryPM25: conc.PrimaryPM25[i],
-			SOA:         conc.SOA[i],
-			TotalPM25:   tpm,
-		}
-		err = o.Encode(r)
-		if err != nil {
-			return err
-		}
-	}
-	o.Close()
-	// Projection definition. This may need to be changed for a different
-	// spatial domain.
-	// TODO: Make this settable by the user, or at least check to make sure it
-	// matches the InMAPProj configuration variable.
-	const proj4 = `PROJCS["Lambert_Conformal_Conic",GEOGCS["GCS_unnamed ellipse",DATUM["D_unknown",SPHEROID["Unknown",6370997,0]],PRIMEM["Greenwich",0],UNIT["Degree",0.017453292519943295]],PROJECTION["Lambert_Conformal_Conic"],PARAMETER["standard_parallel_1",33],PARAMETER["standard_parallel_2",45],PARAMETER["latitude_of_origin",40],PARAMETER["central_meridian",-97],PARAMETER["false_easting",0],PARAMETER["false_northing",0],UNIT["Meter",1]]`
-	// Create .prj file
-	f, err = os.Create(upload.maybeUpload(OutputFile[0:len(OutputFile)-len(filepath.Ext(OutputFile))] + ".prj"))
-	if err != nil {
-		return fmt.Errorf("error creating output prj file: %v", err)
+	if err = r.Output(o, outputVariables, nil, vgsr); err != nil {
+		return err
 	}
-	fmt.Fprint(f, proj4)
-	f.Close()
 
 	if err := upload.uploadOutput(nil); err != nil {
 		return err

inmaputil/sr_test.go

@@ -66,6 +66,8 @@ func TestSRPredict(t *testing.T) {
 	cfg := InitializeConfig()
 	cfg.Set("SR.OutputFile", "../cmd/inmap/testdata/testSR_golden.ncf")
 	cfg.Set("OutputFile", "../cmd/inmap/testdata/output_SRPredict.shp")
+	cfg.Set("OutputVariables", `{"TotalPM25": "PrimaryPM25 + pNH4 + pSO4 + pNO3 + SOA",
+"TotalPopD": "(exp(log(1.078)/10 * TotalPM25) - 1) * TotalPop * allcause / 100000"}`)
 	cfg.Set("EmissionsShapefiles", []string{"../cmd/inmap/testdata/testEmisSR.shp"})
 
 	cfg.Set("config", "../cmd/inmap/configExample.toml")
@@ -80,12 +82,22 @@ func TestSRPredictAboveTop(t *testing.T) {
 	cfg.Set("config", "../cmd/inmap/configExample.toml")
 	cfg.Set("SR.OutputFile", "../cmd/inmap/testdata/testSR_golden.ncf")
 	cfg.Set("OutputFile", "../cmd/inmap/testdata/output_SRPredict.shp")
+	cfg.Set("OutputVariables", `{"PNH4": "pNH4",
+	"PNO3": "pNO3",
+	"PSO4": "pSO4",
+	"SOA": "SOA",
+	"PrimPM25": "PrimaryPM25",
+	"TotalPM25": "PrimaryPM25 + pNH4 + pSO4 + pNO3 + SOA"}`)
 	cfg.Set("EmissionsShapefiles", []string{"../cmd/inmap/testdata/testEmis.shp"})
+	outputVars, err := checkOutputVars(GetStringMapString("OutputVariables", cfg.Viper))
+	if err != nil {
+		t.Fatal(err)
+	}
 	vcfg, err := VarGridConfig(cfg.Viper)
 	if err != nil {
 		t.Fatal(err)
 	}
-	if err := SRPredict(cfg.GetString("EmissionUnits"), cfg.GetString("SR.OutputFile"), cfg.GetString("OutputFile"), cfg.GetStringSlice("EmissionsShapefiles"), vcfg); err != nil {
+	if err := SRPredict(cfg.GetString("EmissionUnits"), cfg.GetString("SR.OutputFile"), cfg.GetString("OutputFile"), outputVars, cfg.GetStringSlice("EmissionsShapefiles"), vcfg); err != nil {
 		t.Fatal(err)
 	}
 }

io.go

@@ -868,7 +868,7 @@ func (d *InMAP) toArray(varName string, layer int, m Mechanism) []float64 {
 			return o
 		}
 		if layer < 0 || c.Layer == layer {
-			o = append(o, c.getValue(varName, d.popIndices, d.mortIndices, m))
+			o = append(o, c.getValue(varName, d.PopIndices, d.mortIndices, m))
 		}
 		c.mutex.RUnlock()
 	}
@@ -882,16 +882,16 @@ func (c *Cell) getValue(varName string, popIndices, mortIndices map[string]int,
 	if err == nil {
 		return v
 	}
-	if polConv, ok := baselinePolLabels[varName]; ok { // Baseline concentrations
+	if i, ok := popIndices[varName]; ok { // Population
+		return c.PopData[i]
+
+	} else if polConv, ok := baselinePolLabels[varName]; ok { // Baseline concentrations
 		var o float64
 		for i, ii := range polConv.index {
 			o += c.CBaseline[ii] * polConv.conversion[i]
 		}
 		return o
 
-	} else if i, ok := popIndices[varName]; ok { // Population
-		return c.PopData[i]
-
 	} else if i, ok := mortIndices[varName]; ok { // Mortality rate
 		return c.MortData[i]
 
@@ -919,11 +919,11 @@ func (d *InMAP) getUnits(varName string, m Mechanism) string {
 	}
 	if _, ok := baselinePolLabels[varName]; ok { // Concentrations
 		return "μg/m³"
-	} else if _, ok := d.popIndices[varName]; ok { // Population
+	} else if _, ok := d.PopIndices[varName]; ok { // Population
 		return "people/grid cell"
 	} else if _, ok := d.mortIndices[varName]; ok { // Mortality Rate
 		return "deaths/100,000"
-	} else if _, ok := d.popIndices[strings.Replace(varName, " deaths", "", 1)]; ok {
+	} else if _, ok := d.PopIndices[strings.Replace(varName, " deaths", "", 1)]; ok {
 		// Mortalities
 		return "deaths/grid cell"
 	}
@@ -964,7 +964,7 @@ func (d *InMAP) OutputOptions(m Mechanism) (names []string, descriptions []strin
 
 	// Population
 	var tempPop []string
-	for pop := range d.popIndices {
+	for pop := range d.PopIndices {
 		tempPop = append(tempPop, pop)
 	}
 	sort.Strings(tempPop)

run.go

@@ -162,7 +162,7 @@ func SteadyStateConvergenceCheck(numIterations int, popGridColumn string, m Mech
 	iteration := 0
 
 	return func(d *InMAP) error {
-		popIndex := d.popIndices[popGridColumn]
+		popIndex := d.PopIndices[popGridColumn]
 
 		if d.Dt == 0 {
 			return fmt.Errorf("inmap: timestep is zero")

save.go

@@ -92,9 +92,9 @@ func Load(r io.Reader, config *VarGridConfig, emis *Emissions, m Mechanism) Doma
 func (d *InMAP) initFromCells(cells []*Cell, emis *Emissions, config *VarGridConfig, m Mechanism) error {
 	d.init()
 	// Create a list of array indices for each population type.
-	d.popIndices = make(map[string]int)
+	d.PopIndices = make(map[string]int)
 	for i, p := range config.CensusPopColumns {
-		d.popIndices[p] = i
+		d.PopIndices[p] = i
 	}
 	d.mortIndices = make(map[string]int)
 	mortRateColumns := make([]string, len(config.MortalityRateColumns))

sr/sr.go

@@ -35,8 +35,8 @@ import (
 	"sync"
 	"time"
 
-	"github.com/ctessum/cdf"
 	"github.com/cenkalti/backoff"
+	"github.com/ctessum/cdf"
 	"github.com/ctessum/geom"
 	"github.com/ctessum/geom/encoding/shp"
 	"github.com/lnashier/viper"

sr/sr_test.go

@@ -28,8 +28,8 @@ import (
 	"strings"
 	"testing"
 
-	"github.com/ctessum/cdf"
 	"github.com/BurntSushi/toml"
+	"github.com/ctessum/cdf"
 	"github.com/gonum/floats"
 	"github.com/spatialmodel/inmap"
 	"github.com/spatialmodel/inmap/cloud"