Update rail.py

rail/rail.py

@@ -1,3 +1,4 @@
+import copy
 import numpy as np
 import pandas as pd
 from matplotlib import pyplot as plt
@@ -72,22 +73,24 @@ def new(self, name: str, threat_source: ThreatSource) -> ThreatEvent:
 
 
 class Control(UserDict):
-    def __init__(self, name: str, cost: float, likelihood_reduction: float) -> None:
+    def __init__(self, name: str, cost: float, likelihood_reduction: float, implemented: bool=True) -> None:
         self.data = {}
         self.data['name'] = name
         self.data['cost'] = cost
         self.data['likelihood_reduction'] = likelihood_reduction
+        self.data['implemented'] = implemented
 
 
 class Controls(UserDict):
     def __init__(self) -> None:
         self.data = {}
-        self.costs = 0.0
 
     def new(self, name: str, cost: float, likelihood_reduction: float) -> Control:
         self.data[name] = Control(name, cost, likelihood_reduction)
-        self.costs += cost
-        return self.data[name]      
+        return self.data[name]     
+
+    def costs(self):
+        return np.sum(list(map(lambda x: x['cost'] if x['implemented'] == True else 0, self.data.values())))
 
 
 class Vulnerability(UserDict):
@@ -97,7 +100,6 @@ def __init__(self, threat_event: ThreatEvent, system: Tree, controls: [Control]
         self.data['threat_event'] = threat_event
         self.data['system'] = system
         self.data['controls'] = controls
-        self.data['likelihood_reduction'] = np.product(list(map(lambda x: x['likelihood_reduction'], controls)))
 
 
 class Vulnerabilities(UserDict):
@@ -166,16 +168,20 @@ def __init__(self, vulnerability: Vulnerability, likelihood: Likelihood, impact:
         self.data['vulnerability'] = vulnerability
         self.data['likelihood'] = likelihood
         self.data['impact'] = impact
-        self.data['mean'] = impact['mean'] * likelihood['lam'] * vulnerability['likelihood_reduction']
 
+    def evaluate_deterministic(self) -> float:
+        likelihood_reduction = np.product(list(map(lambda x: x['likelihood_reduction'] if x['implemented'] == True else 1, self.data['vulnerability']['controls'])))
+        return self.data['likelihood']['lam'] * self.data['impact']['mean'] * likelihood_reduction
     def evaluate_lognormal(self, iterations: int = 1000) -> float:
-        return lognorm.ppf(np.random.rand(iterations), s=self.data['impact']['sigma'], scale=np.exp(self.data['impact']['mu'])) * np.random.poisson(lam=self.data['likelihood']['lam'] * self.data['vulnerability']['likelihood_reduction'], size=iterations)
+        likelihood_reduction = np.product(list(map(lambda x: x['likelihood_reduction'] if x['implemented'] == True else 1, self.data['vulnerability']['controls'])))
+        return lognorm.ppf(np.random.rand(iterations), s=self.data['impact']['sigma'], scale=np.exp(self.data['impact']['mu'])) * np.random.poisson(lam=self.data['likelihood']['lam'] * likelihood_reduction, size=iterations)
 
 
 class Risks(UserDict):
     def __init__(self) -> None:
         self.data = {}
-        self.dataframe = pd.DataFrame(columns=['Threat Source', 'Threat Event', 'System', 'Controls', 'Impact', 'Impact (mean)', 'Likelihood (mean)', 'Expected Loss (mean)'])
+        self.dataframe = pd.DataFrame(columns=['Threat Source', 'Threat Event', 'System', 'Controls', 'Impact', 'Impact (mean)', 'Likelihood (mean)'])
+        self.cost_loss = []
 
     def new(self, vulnerability: Vulnerability, likelihood: Likelihood, impact: Impact) -> Risk:
         name = vulnerability['name'] + ' -> ' + likelihood['name'] + ' -> ' + impact['name']
@@ -188,22 +194,65 @@ def new(self, vulnerability: Vulnerability, likelihood: Likelihood, impact: Impa
             'Impact': impact['name'],
             'Impact (mean)': impact['mean'],
             'Likelihood (mean)': likelihood['lam'],
-            'Expected Loss (mean)': impact['mean'] * likelihood['lam'] * vulnerability['likelihood_reduction'],
         }, ignore_index=True)
         return self.data[name]
 
     def calculate_stochastic_risks(self, interations: int = 100000):
         risks_scores = np.array(list(map(lambda row: row.evaluate_lognormal(interations), self.data.values())))
         return risks_scores.sum(axis=0)
 
-    def plot(self, range=(0, 1000000)):
+    def plot(self, axes=None):
         plt.title('expected loss')
         plt.xlabel('loss')
         plt.ylabel('probability')
-        return plt.hist(self.calculate_stochastic_risks(), histtype='stepfilled', bins=100, cumulative=-1, normed=True, range=range)
+        return plt.hist(self.calculate_stochastic_risks(), histtype='step', bins=10000, cumulative=-1, normed=True, axes=axes)
 
     def expected_loss_stochastic_mean(self, interations: int = 1000) -> float:
         return self.calculate_stochastic_risks(interations).sum() / interations
 
     def expected_loss_deterministic_mean(self) -> float:
-        return self.dataframe['Expected Loss (mean)'].sum()
+        return np.array(list(map(lambda row: row.evaluate_deterministic(), self.data.values()))).sum()
+
+    def calculate_dataframe_deterministic_mean(self):
+        df = risks.dataframe.copy()
+        df['Risk (mean)'] = list(map(lambda x: x.evaluate_deterministic(), self.data.values()))
+        return df
+
+    def determine_optimum_controls(self, controls, controls_to_optimize):
+        if not controls_to_optimize:
+            loss = self.expected_loss_deterministic_mean()
+            cost = controls.costs()
+            self.cost_loss.append({'cost': cost, 'loss': loss})
+            return {'loss': loss, 'cost': cost, 'controls': copy.deepcopy(controls)}
+        else:
+            controls_to_optimize_new_list = list(controls_to_optimize)
+            control = controls_to_optimize_new_list.pop()
+            controls[control]['implemented'] = False
+            control_off = self.determine_optimum_controls(controls, controls_to_optimize_new_list)
+            controls[control]['implemented'] = True
+            control_on = self.determine_optimum_controls(controls, controls_to_optimize_new_list)
+            if control_on['loss'] + control_on['cost'] < control_off['loss'] + control_off['cost']:
+                optimal_control = control_on
+            else:
+                optimal_control = control_off
+            return optimal_control
+
+    def set_optimum_controls(self, controls):
+        optimum_controls = self.determine_optimum_controls(controls, controls)
+        for control in optimum_controls['controls']:
+            if optimum_controls['controls'][control]['implemented'] == True:
+                controls[control]['implemented'] = True
+            else:
+                controls[control]['implemented'] = False
+        df = pd.DataFrame(list(optimum_controls['controls'].values())).set_index('name')
+        return df
+
+    def plot_risk_cost_matrix(self, controls=controls, axes=None):
+        self.set_optimum_controls(controls)
+        df = pd.DataFrame(self.cost_loss)
+        plt.title('residual risk versus control cost')
+        plt.ylabel('residual risk')
+        plt.xlabel('control cost')
+        plt.scatter(df['cost'], df['loss'], axes=axes)
+        plt.scatter(controls.costs(), risks.expected_loss_deterministic_mean(), color='red', axes=axes)
+        axes.set_xlim(xmin=0)