forked from microsoft/SynapseML
-
Notifications
You must be signed in to change notification settings - Fork 0
/
DoubleMLModel.py
26 lines (20 loc) · 828 Bytes
/
DoubleMLModel.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
# Copyright (C) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See LICENSE in project root for information.
import sys
if sys.version >= "3":
basestring = str
from synapse.ml.causal._DoubleMLModel import _DoubleMLModel
from pyspark.ml.common import inherit_doc
import numpy as np
@inherit_doc
class DoubleMLModel(_DoubleMLModel):
def getAvgTreatmentEffect(self):
return sum(self.getRawTreatmentEffects()) / len(self.getRawTreatmentEffects())
def getConfidenceInterval(self):
ciLowerBound = np.percentile(
self.getRawTreatmentEffects(), 100 * (1 - self.getConfidenceLevel())
)
ciUpperBound = np.percentile(
self.getRawTreatmentEffects(), self.getConfidenceLevel() * 100
)
return [ciLowerBound, ciUpperBound]