feat: added test case for custom callbacks. works great and somehow serializes.

dcrescim · dcrescim · commit 7fa5c4259902 · 2022-05-18T21:41:05.000-07:00
diff --git a/src/linear_model/LinearRegression.test.ts b/src/linear_model/LinearRegression.test.ts
@@ -17,6 +17,38 @@ describe('LinearRegression', function () {
     expect(roughlyEqual(lr.intercept as number, 0)).toBe(true)
   }, 30000)
 
+  it('Works on arrays (small example) with custom callbacks', async function () {
+    let trainingHasStarted = false
+    const onTrainBegin = async (logs: any) => {
+      trainingHasStarted = true
+      console.log('training begins')
+    }
+    const lr = new LinearRegression({
+      modelFitOptions: { callbacks: [new tf.CustomCallback({ onTrainBegin })] }
+    })
+    await lr.fit([[1], [2]], [2, 4])
+    expect(tensorEqual(lr.coef, tf.tensor1d([2]), 0.1)).toBe(true)
+    expect(roughlyEqual(lr.intercept as number, 0)).toBe(true)
+    expect(trainingHasStarted).toBe(true)
+  }, 30000)
+
+  it('Works on arrays (small example) with custom callbacks', async function () {
+    let trainingHasStarted = false
+    const onTrainBegin = async (logs: any) => {
+      trainingHasStarted = true
+      console.log('training begins')
+    }
+    const lr = new LinearRegression({
+      modelFitOptions: { callbacks: [new tf.CustomCallback({ onTrainBegin })] }
+    })
+    await lr.fit([[1], [2]], [2, 4])
+
+    const serialized = await lr.toJSON()
+    const newModel = await fromJSON(serialized)
+    expect(tensorEqual(newModel.coef, tf.tensor1d([2]), 0.1)).toBe(true)
+    expect(roughlyEqual(newModel.intercept as number, 0)).toBe(true)
+  }, 30000)
+
   it('Works on small multi-output example (small example)', async function () {
     const lr = new LinearRegression()
     await lr.fit(
diff --git a/src/linear_model/LinearRegression.ts b/src/linear_model/LinearRegression.ts
@@ -41,7 +41,6 @@ export interface LinearRegressionParams {
    */
   fitIntercept?: boolean
   modelFitOptions?: Partial<ModelFitArgs>
-
 }
 
 /*
@@ -53,7 +52,7 @@ Next steps:
 /** Linear Least Squares
  * @example
  * ```js
- * import {LinearRegression} from 'scikitjs'
+ * import { LinearRegression } from 'scikitjs'
  *
  * let X = [
  *  [1, 2],
@@ -63,13 +62,16 @@ Next steps:
  *  [10, 20]
  * ]
  * let y = [3, 5, 8, 8, 30]
- * const lr = new LinearRegression({fitIntercept: false})
+ * const lr = new LinearRegression({ fitIntercept: false })
   await lr.fit(X, y)
   lr.coef.print() // probably around [1, 1]
  * ```
  */
 export class LinearRegression extends SGDRegressor {
-  constructor({ fitIntercept = true, modelFitOptions }: LinearRegressionParams = {}) {
+  constructor({
+    fitIntercept = true,
+    modelFitOptions
+  }: LinearRegressionParams = {}) {
     let tf = getBackend()
     super({
       modelCompileArgs: {
