GH-15809: add dispersion_estimated parameter to GLM mojo

h2oai · Jan 26, 2024 · 49ba70a · 49ba70a
1 parent 3a45e5c
commit 49ba70a
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 0 deletions.
diff --git a/h2o-algos/src/main/java/hex/glm/GLMMojoWriter.java b/h2o-algos/src/main/java/hex/glm/GLMMojoWriter.java
@@ -39,6 +39,8 @@ protected void writeModelData() throws IOException {
 
     if (GLMModel.GLMParameters.Family.tweedie.equals(model._parms._family))
       writekv("tweedie_link_power", model._parms._tweedie_link_power);
+
+    writekv("dispersion_estimated", model._parms._compute_p_values ? model._parms._dispersion_estimated : 0);
   }
 
 }
diff --git a/h2o-genmodel/src/main/java/hex/genmodel/algos/glm/GlmMojoModelBase.java b/h2o-genmodel/src/main/java/hex/genmodel/algos/glm/GlmMojoModelBase.java
@@ -18,6 +18,8 @@ public abstract class GlmMojoModelBase extends MojoModel {
 
   String _family;
   boolean _versionSupportOffset;
+
+  double _dispersion_estimated;
 
   GlmMojoModelBase(String[] columns, String[][] domains, String responseColumn) {
     super(columns, domains, responseColumn);
@@ -26,6 +28,10 @@ public abstract class GlmMojoModelBase extends MojoModel {
   public double[] getBeta() {
     return _beta;
   }
+
+  public double getDispersionEstimated() {
+    return _dispersion_estimated;
+  }
 
   void init() {
     _versionSupportOffset = _mojo_version >= 1.1;

diff --git a/h2o-genmodel/src/main/java/hex/genmodel/algos/glm/GlmMojoReader.java b/h2o-genmodel/src/main/java/hex/genmodel/algos/glm/GlmMojoReader.java
@@ -29,6 +29,7 @@ protected void readModelData() throws IOException {
 
     _model._beta = readkv("beta");
     _model._family = readkv("family");
+    _model._dispersion_estimated = readkv("dispersion_estimated");
 
     if (_model instanceof GlmMojoModel) {
       GlmMojoModel m = (GlmMojoModel) _model;