text: import sys
  import unittest
  import traceback
  
  import numpy as np
  import openmdao.api as om
  import numpy.testing as npt
  from wisdem.drivetrainse.drivetrain import DrivetrainSE
  
  npts = 12
  
  
  def set_common(prob, opt):
      prob["n_blades"] = 3
      prob["rotor_diameter"] = 120.0
      prob["machine_rating"] = 5e3
      prob["D_top"] = 6.5
      prob["rated_torque"] = 10.25e6  # rev 1 9.94718e6
  
      prob["F_hub"] = np.array([2409.750e3, 0.0, 74.3529e2]).reshape((3, 1))
      prob["M_hub"] = np.array([-1.83291e4, 6171.7324e2, 5785.82946e2]).reshape((3, 1))
  
      prob["E_mat"] = 210e9 * np.ones((1, 3))
      prob["G_mat"] = 80.8e9 * np.ones((1, 3))
      prob["Xt_mat"] = 1e7 * np.ones((1, 3))
      prob["rho_mat"] = 7850.0 * np.ones(1)
      prob["sigma_y_mat"] = 250e6 * np.ones(1)
      prob["unit_cost_mat"] = 3.0 * np.ones(1)
      prob["lss_material"] = prob["hss_material"] = prob["hub_material"] = prob["spinner_material"] = prob[
          "bedplate_material"
      ] = "steel"
      prob["material_names"] = ["steel"]
  
      prob["blade_mass"] = 17000.0
      prob["pitch_system.BRFM"] = 1.0e6
      prob["pitch_system_scaling_factor"] = 0.54
  
      prob["blade_root_diameter"] = 4.0
      prob["flange_t2shell_t"] = 4.0
      prob["flange_OD2hub_D"] = 0.5
      prob["flange_ID2flange_OD"] = 0.8
      prob["hub_rho"] = 7200.0
      prob["hub_Xy"] = 200.0e6
      prob["hub_stress_concentration"] = 2.5
  
      prob["n_front_brackets"] = 3
      prob["n_rear_brackets"] = 3
      prob["clearance_hub_spinner"] = 0.5
      prob["spin_hole_incr"] = 1.2
      prob["spinner_gust_ws"] = 70.0
  
      prob["hub_diameter"] = 6.235382907247958
      prob["minimum_rpm"] = 2  # 8.68                # rpm 9.6
      prob["rated_rpm"] = 10  # 8.68                # rpm 9.6
      prob["blades_I"] = 1e3 * np.ones(6)
      prob["blades_I"][0] = 199200777.51 * 30.0 * 5 / 10.0 / np.pi
  
      prob["bear1.bearing_type"] = "CARB"
      prob["bear2.bearing_type"] = "SRB"
  
      return prob
  
  
  class TestGroup(unittest.TestCase):
      def testDirectDrive_withGen(self):
  
          opt = {}
          opt["WISDEM"] = {}
          opt["WISDEM"]["DriveSE"] = {}
          opt["WISDEM"]["DriveSE"]["direct"] = True
          opt["WISDEM"]["DriveSE"]["hub"] = {}
          opt["WISDEM"]["DriveSE"]["hub"]["hub_gamma"] = 2.0
          opt["WISDEM"]["DriveSE"]["hub"]["spinner_gamma"] = 1.5
          opt["WISDEM"]["DriveSE"]["gamma_f"] = 1.35
          opt["WISDEM"]["DriveSE"]["gamma_m"] = 1.3
          opt["WISDEM"]["DriveSE"]["gamma_n"] = 1.0
          opt["WISDEM"]["RotorSE"] = {}
          opt["WISDEM"]["RotorSE"]["n_pc"] = 20
          opt["materials"] = {}
          opt["materials"]["n_mat"] = 1
          opt["WISDEM"]["GeneratorSE"] = {}
          opt["WISDEM"]["GeneratorSE"]["type"] = "pmsg_outer"
          opt["flags"] = {}
          opt["flags"]["generator"] = True
  
          prob = om.Problem()
          prob.model = DrivetrainSE(modeling_options=opt, n_dlcs=1)
          prob.setup()
          prob = set_common(prob, opt)
  
          prob["upwind"] = True
  
          prob["L_12"] = 2.0
          prob["L_h1"] = 1.0
          prob["L_generator"] = 3.25
          prob["overhang"] = 6.25
          prob["drive_height"] = 4.875
          prob["tilt"] = 4.0
          prob["access_diameter"] = 0.9
  
          myones = np.ones(2)
          prob["lss_diameter"] = 3.3 * myones
          prob["nose_diameter"] = 2.2 * myones
          prob["lss_wall_thickness"] = 0.45 * myones
          prob["nose_wall_thickness"] = 0.1 * myones
          prob["bedplate_wall_thickness"] = 0.06 * np.ones(4)
          prob["bear1.D_shaft"] = 2.2
          prob["bear2.D_shaft"] = 2.2
          prob["generator.D_shaft"] = 3.3
          prob["generator.D_nose"] = 2.2
  
          prob["generator.P_mech"] = 10.71947704e6  # rev 1 9.94718e6
          prob["generator.rad_ag"] = 4.0  # rev 1  4.92
          prob["generator.len_s"] = 1.7  # rev 2.3
          prob["generator.h_s"] = 0.7  # rev 1 0.3
          prob["generator.p"] = 70  # 100.0    # rev 1 160
          prob["generator.h_m"] = 0.005  # rev 1 0.034
          prob["generator.h_ys"] = 0.04  # rev 1 0.045
          prob["generator.h_yr"] = 0.06  # rev 1 0.045
          prob["generator.b"] = 2.0
          prob["generator.c"] = 5.0
          prob["generator.B_tmax"] = 1.9
          prob["generator.E_p"] = 3300 / np.sqrt(3)
          prob["generator.t_r"] = 0.05  # Rotor disc thickness
          prob["generator.h_sr"] = 0.04  # Rotor cylinder thickness
          prob["generator.t_s"] = 0.053  # Stator disc thickness
          prob["generator.h_ss"] = 0.04  # Stator cylinder thickness
          prob["generator.u_allow_pcent"] = 8.5  # % radial deflection
          prob["generator.y_allow_pcent"] = 1.0  # % axial deflection
          prob["generator.z_allow_deg"] = 0.05  # torsional twist
          prob["generator.sigma"] = 60.0e3  # Shear stress
          prob["generator.B_r"] = 1.279
          prob["generator.ratio_mw2pp"] = 0.8
          prob["generator.h_0"] = 5e-3
          prob["generator.h_w"] = 4e-3
          prob["generator.k_fes"] = 0.8
          prob["generator.C_Cu"] = 4.786  # Unit cost of Copper $/kg
          prob["generator.C_Fe"] = 0.556  # Unit cost of Iron $/kg
          prob["generator.C_Fes"] = 0.50139  # specific cost of Structural_mass $/kg
          prob["generator.C_PM"] = 95.0
          prob["generator.rho_Fe"] = 7700.0  # Steel density Kg/m3
          prob["generator.rho_Fes"] = 7850  # structural Steel density Kg/m3
          prob["generator.rho_Copper"] = 8900.0  # copper density Kg/m3
          prob["generator.rho_PM"] = 7450.0  # typical density Kg/m3 of neodymium magnets
  
          try:
              prob.run_model()
              self.assertTrue(True)
          except Exception:
              traceback.print_exc(file=sys.stdout)
              self.assertTrue(False)
  
          # Test that key outputs are filled
          self.assertGreater(prob["nacelle_mass"], 100e3)
          self.assertLess(prob["nacelle_cm"][0], 0.0)
          self.assertGreater(prob["nacelle_cm"][2], 0.0)
          self.assertGreater(prob["rna_mass"], 100e3)
          self.assertLess(prob["rna_cm"][0], 0.0)
          self.assertGreater(prob["rna_cm"][2], 0.0)
          self.assertGreater(prob["generator_mass"], 100e3)
          self.assertGreater(prob["generator_cost"], 100e3)
          npt.assert_array_less(100e3, prob["generator_I"])
          npt.assert_array_less(0.8, prob["generator_efficiency"])
          npt.assert_array_less(prob["generator_efficiency"], 1.0)
  
      def testDirectDrive_withSimpleGen(self):
  
          opt = {}
          opt["WISDEM"] = {}
          opt["WISDEM"]["DriveSE"] = {}
          opt["WISDEM"]["DriveSE"]["direct"] = True
          opt["WISDEM"]["DriveSE"]["hub"] = {}
          opt["WISDEM"]["DriveSE"]["hub"]["hub_gamma"] = 2.0
          opt["WISDEM"]["DriveSE"]["hub"]["spinner_gamma"] = 1.5
          opt["WISDEM"]["DriveSE"]["gamma_f"] = 1.35
          opt["WISDEM"]["DriveSE"]["gamma_m"] = 1.3
          opt["WISDEM"]["DriveSE"]["gamma_n"] = 1.0
          opt["WISDEM"]["RotorSE"] = {}
          opt["WISDEM"]["RotorSE"]["n_pc"] = 20
          opt["materials"] = {}
          opt["materials"]["n_mat"] = 1
          opt["flags"] = {}
          opt["flags"]["generator"] = False
  
          prob = om.Problem()
          prob.model = DrivetrainSE(modeling_options=opt, n_dlcs=1)
          prob.setup()
          prob = set_common(prob, opt)
  
          prob["upwind"] = True
  
          prob["L_12"] = 2.0
          prob["L_h1"] = 1.0
          prob["L_generator"] = 3.25
          prob["overhang"] = 6.25
          prob["drive_height"] = 4.875
          prob["tilt"] = 4.0
          prob["access_diameter"] = 0.9
  
          myones = np.ones(2)
          prob["lss_diameter"] = 3.3 * myones
          prob["nose_diameter"] = 2.2 * myones
          prob["lss_wall_thickness"] = 0.45 * myones
          prob["nose_wall_thickness"] = 0.1 * myones
          prob["bedplate_wall_thickness"] = 0.06 * np.ones(4)
          prob["bear1.D_shaft"] = 2.2
          prob["bear2.D_shaft"] = 2.2
  
          prob["generator_efficiency_user"] = np.zeros((20, 2))
  
          try:
              prob.run_model()
              self.assertTrue(True)
          except Exception:
              traceback.print_exc(file=sys.stdout)
              self.assertTrue(False)
  
          # Test that key outputs are filled
          self.assertGreater(prob["nacelle_mass"], 100e3)
          self.assertLess(prob["nacelle_cm"][0], 0.0)
          self.assertGreater(prob["nacelle_cm"][2], 0.0)
          self.assertGreater(prob["rna_mass"], 100e3)
          self.assertLess(prob["rna_cm"][0], 0.0)
          self.assertGreater(prob["rna_cm"][2], 0.0)
          self.assertGreater(prob["generator_mass"], 100e3)
          # self.assertGreater(prob["generator_cost"], 100e3)
          npt.assert_array_less(100e3, prob["generator_I"])
          npt.assert_array_less(0.8, prob["generator_efficiency"])
          npt.assert_array_less(prob["generator_efficiency"], 1.0)
  
      def testGeared_withGen(self):
  
          opt = {}
          opt["WISDEM"] = {}
          opt["WISDEM"]["DriveSE"] = {}
          opt["WISDEM"]["DriveSE"]["direct"] = False
          opt["WISDEM"]["DriveSE"]["hub"] = {}
          opt["WISDEM"]["DriveSE"]["hub"]["hub_gamma"] = 2.0
          opt["WISDEM"]["DriveSE"]["hub"]["spinner_gamma"] = 1.5
          opt["WISDEM"]["DriveSE"]["gamma_f"] = 1.35
          opt["WISDEM"]["DriveSE"]["gamma_m"] = 1.3
          opt["WISDEM"]["DriveSE"]["gamma_n"] = 1.0
          opt["WISDEM"]["GeneratorSE"] = {}
          opt["WISDEM"]["GeneratorSE"]["type"] = "dfig"
          opt["WISDEM"]["RotorSE"] = {}
          opt["WISDEM"]["RotorSE"]["n_pc"] = 20
          opt["materials"] = {}
          opt["materials"]["n_mat"] = 1
          opt["flags"] = {}
          opt["flags"]["generator"] = True
  
          prob = om.Problem()
          prob.model = DrivetrainSE(modeling_options=opt, n_dlcs=1)
          prob.setup()
          prob = set_common(prob, opt)
  
          prob["upwind"] = True
  
          prob["L_12"] = 2.0
          prob["L_h1"] = 1.0
          prob["overhang"] = 2.0
          prob["drive_height"] = 4.875
          prob["L_hss"] = 1.5
          prob["L_generator"] = 1.25
          prob["L_gearbox"] = 1.1
          prob["tilt"] = 5.0
  
          myones = np.ones(2)
          prob["lss_diameter"] = 2.3 * myones
          prob["lss_wall_thickness"] = 0.05 * myones
          prob["hss_diameter"] = 2.0 * myones
          prob["hss_wall_thickness"] = 0.05 * myones
  
          prob["bedplate_flange_width"] = 1.5
          prob["bedplate_flange_thickness"] = 0.05
          # prob['bedplate_web_height'] = 1.0
          prob["bedplate_web_thickness"] = 0.05
          prob["bear1.D_shaft"] = 2.3
          prob["bear2.D_shaft"] = 2.3
  
          prob["planet_numbers"] = np.array([3, 3, 0])
          prob["gear_configuration"] = "eep"
          # prob['shaft_factor'] = 'normal'
          prob["gear_ratio"] = 90.0
  
          prob["generator.rho_Fe"] = 7700.0
          prob["generator.rho_Fes"] = 7850.0
          prob["generator.rho_Copper"] = 8900.0
          prob["generator.rho_PM"] = 7450.0
          prob["generator.B_r"] = 1.2
          prob["generator.E"] = 2e11
          prob["generator.G"] = 79.3e9
          prob["generator.P_Fe0e"] = 1.0
          prob["generator.P_Fe0h"] = 4.0
          prob["generator.S_N"] = -0.002
          prob["generator.alpha_p"] = 0.5 * np.pi * 0.7
          prob["generator.b_r_tau_r"] = 0.45
          prob["generator.b_ro"] = 0.004
          prob["generator.b_s_tau_s"] = 0.45
          prob["generator.b_so"] = 0.004
          prob["generator.cofi"] = 0.85
          prob["generator.freq"] = 60
          prob["generator.h_i"] = 0.001
          prob["generator.h_sy0"] = 0.0
          prob["generator.h_w"] = 0.005
          prob["generator.k_fes"] = 0.9
          prob["generator.k_s"] = 0.2
          prob["generator.m"] = 3
          prob["generator.mu_0"] = np.pi * 4e-7
          prob["generator.mu_r"] = 1.06
          prob["generator.p"] = 3.0
          prob["generator.phi"] = np.deg2rad(90)
          prob["generator.ratio_mw2pp"] = 0.7
          prob["generator.resist_Cu"] = 1.8e-8 * 1.4
          prob["generator.sigma"] = 40e3
          prob["generator.v"] = 0.3
          prob["generator.y_tau_p"] = 1.0
          prob["generator.y_tau_pr"] = 10.0 / 12
          prob["generator.cofi"] = 0.9
          prob["generator.y_tau_p"] = 12.0 / 15.0
          prob["generator.sigma"] = 21.5e3
          prob["generator.rad_ag"] = 0.61
          prob["generator.len_s"] = 0.49
          prob["generator.h_s"] = 0.08
          prob["generator.I_0"] = 40.0
          prob["generator.B_symax"] = 1.3
          prob["generator.S_Nmax"] = -0.2
          prob["generator.h_0"] = 0.01
          prob["generator.k_fillr"] = 0.55
          prob["generator.k_fills"] = 0.65
          prob["generator.q1"] = 5
          prob["generator.q2"] = 4
  
          # prob['generator.D_shaft'] = 2.3
  
          try:
              prob.run_model()
              self.assertTrue(True)
          except Exception:
              traceback.print_exc(file=sys.stdout)
              self.assertTrue(False)
  
          # Test that key outputs are filled
          self.assertGreater(prob["nacelle_mass"], 100e3)
          self.assertGreater(prob["nacelle_cm"][0], 0.0)
          self.assertGreater(prob["nacelle_cm"][2], 0.0)
          self.assertGreater(prob["rna_mass"], 100e3)
          self.assertGreater(prob["rna_cm"][0], 0.0)
          self.assertGreater(prob["rna_cm"][2], 0.0)
          self.assertGreater(prob["generator_mass"], 10e3)
          self.assertGreater(prob["generator_cost"], 10e3)
          npt.assert_array_less(1e3, prob["generator_I"])
          npt.assert_array_less(0.2, prob["generator_efficiency"][1:])
          npt.assert_array_less(prob["generator_efficiency"], 1.0)
  
      def testGeared_withSimpleGen(self):
  
          opt = {}
          opt["WISDEM"] = {}
          opt["WISDEM"]["DriveSE"] = {}
          opt["WISDEM"]["DriveSE"]["direct"] = False
          opt["WISDEM"]["DriveSE"]["hub"] = {}
          opt["WISDEM"]["DriveSE"]["hub"]["hub_gamma"] = 2.0
          opt["WISDEM"]["DriveSE"]["hub"]["spinner_gamma"] = 1.5
          opt["WISDEM"]["DriveSE"]["gamma_f"] = 1.35
          opt["WISDEM"]["DriveSE"]["gamma_m"] = 1.3
          opt["WISDEM"]["DriveSE"]["gamma_n"] = 1.0
          opt["flags"] = {}
          opt["flags"]["generator"] = False
          opt["WISDEM"]["RotorSE"] = {}
          opt["WISDEM"]["RotorSE"]["n_pc"] = 20
          opt["materials"] = {}
          opt["materials"]["n_mat"] = 1
  
          prob = om.Problem()
          prob.model = DrivetrainSE(modeling_options=opt, n_dlcs=1)
          prob.setup()
          prob = set_common(prob, opt)
  
          prob["upwind"] = True
  
          prob["L_12"] = 2.0
          prob["L_h1"] = 1.0
          prob["overhang"] = 2.0
          prob["drive_height"] = 4.875
          prob["L_hss"] = 1.5
          prob["L_generator"] = 1.25
          prob["L_gearbox"] = 1.1
          prob["tilt"] = 5.0
  
          myones = np.ones(2)
          prob["lss_diameter"] = 2.3 * myones
          prob["lss_wall_thickness"] = 0.05 * myones
          prob["hss_diameter"] = 2.0 * myones
          prob["hss_wall_thickness"] = 0.05 * myones
  
          prob["bedplate_flange_width"] = 1.5
          prob["bedplate_flange_thickness"] = 0.05
          # prob['bedplate_web_height'] = 1.0
          prob["bedplate_web_thickness"] = 0.05
          prob["bear1.D_shaft"] = 2.3
          prob["bear2.D_shaft"] = 2.3
          # prob['generator.D_shaft'] = 2.3
  
          prob["planet_numbers"] = np.array([3, 3, 0])
          prob["gear_configuration"] = "eep"
          # prob['shaft_factor'] = 'normal'
          prob["gear_ratio"] = 90.0
  
          prob["generator_efficiency_user"] = np.zeros((20, 2))
  
          try:
              prob.run_model()
              self.assertTrue(True)
          except Exception:
              traceback.print_exc(file=sys.stdout)
              self.assertTrue(False)
  
          # Test that key outputs are filled
          self.assertGreater(prob["nacelle_mass"], 100e3)
          self.assertGreater(prob["nacelle_cm"][0], 0.0)
          self.assertGreater(prob["nacelle_cm"][2], 0.0)
          self.assertGreater(prob["rna_mass"], 100e3)
          self.assertGreater(prob["rna_cm"][0], 0.0)
          self.assertGreater(prob["rna_cm"][2], 0.0)
          self.assertGreater(prob["generator_mass"], 10e3)
          # self.assertGreater(prob["generator_cost"], 10e3)
          npt.assert_array_less(1e3, prob["generator_I"])
          npt.assert_array_less(0.8, prob["generator_efficiency"])
          npt.assert_array_less(prob["generator_efficiency"], 1.0)
  
  
  def suite():
      suite = unittest.TestSuite()
      suite.addTest(unittest.makeSuite(TestGroup))
      return suite
  
  
  if __name__ == "__main__":
      result = unittest.TextTestRunner().run(suite())
  
      if result.wasSuccessful():
          exit(0)
      else:
          exit(1)
  