code: def partially_linear_dataset(
      beta,
      num_common_causes,
      num_unobserved_common_causes=0,
      strength_unobserved_confounding=1,
      num_samples=500,
      num_treatments=None,
      treatment_is_binary=True,
      treatment_is_category=False,
      outcome_is_binary=False,
      stochastic_discretization=True,
      num_discrete_common_causes=0,
      stddev_treatment_noise=1,
      stddev_outcome_noise=0,
      one_hot_encode=False,
      training_sample_size=10,
      random_state=0,
  ):
      assert not (treatment_is_binary and treatment_is_category)
      num_outcomes = 1
      beta = np.array(beta, dtype=np.float64, ndmin=1)
      if num_treatments is None:
          num_treatments = beta.size
      else:
          beta = np.resize(beta, num_treatments)
      num_cont_common_causes = (
          num_common_causes - num_discrete_common_causes
      )
  
      if num_common_causes > 0:
          range_c1 = 0.5 + np.max(np.absolute(beta)) * 0.5
          means = np.random.uniform(-1, 1, num_common_causes)
          cov_mat = np.diag(np.ones(num_common_causes))
          W = np.random.multivariate_normal(means, cov_mat, num_samples)
          W_with_dummy = convert_to_categorical(
              W,
              num_common_causes,
              num_discrete_common_causes,
              quantiles=[0.25, 0.5, 0.75],
              one_hot_encode=one_hot_encode,
          )
          c1 = np.random.uniform(
              0, range_c1, (W_with_dummy.shape[1], num_treatments)
          )
          # assuming that all unobserved common causes are numerical and are not affected by one hot encoding
          if num_unobserved_common_causes > 0:
              c1[:num_unobserved_common_causes] = (
                  c1[:num_unobserved_common_causes]
                  * strength_unobserved_confounding
              )
              for i in range(
                  num_unobserved_common_causes, W_with_dummy.shape[1]
              ):
                  c1[i] = (
                      c1[i] * strength_unobserved_confounding / (2**i)
                  )
  
      # Creating a NN to simulate the nuisance function
      hidden_layer_arch = (50, 50, 50)
      neural_network = MLPRegressor(
          random_state=random_state,
          hidden_layer_sizes=hidden_layer_arch,
      )
      x = np.random.randn(training_sample_size, W_with_dummy.shape[1])
      y = np.random.randn(training_sample_size, num_outcomes)
      neural_network.fit(x, np.ravel(y))
  
      T = np.random.normal(
          0, stddev_treatment_noise, (num_samples, num_treatments)
      )
      T += W_with_dummy @ c1
      if treatment_is_binary:
          T = np.vectorize(convert_to_binary)(x=T)
  
      # strength of unobserved confounding
      strength_vec = np.ones(W_with_dummy.shape[1])
      strength_vec[:num_unobserved_common_causes] = (
          strength_vec[:num_unobserved_common_causes]
          * strength_unobserved_confounding
      )
      for i in range(
          num_unobserved_common_causes, W_with_dummy.shape[1]
      ):
          strength_vec[i] = (
              strength_vec[i]
              * strength_unobserved_confounding
              / (2**i)
          )
  
      def _compute_y(T, W, beta, nn, stddev_outcome_noise):
          f_x = nn.predict(W)
          y = np.random.normal(0, stddev_outcome_noise, num_samples)
          y += T @ beta
          y += f_x
          if outcome_is_binary:
              y = np.vectorize(convert_to_binary)(
                  y, stochastic_discretization
              )
          return y
  
      Y = _compute_y(
          T,
          W_with_dummy * strength_vec,
          beta,
          neural_network,
          stddev_outcome_noise,
      )
  
      data = np.column_stack((T, Y))
      if num_common_causes > 0:
          data = np.column_stack((W_with_dummy, data))
  
      # Computing ATE
      T1 = np.ones((num_samples, num_treatments))
      T0 = np.zeros((num_samples, num_treatments))
      ate = np.mean(
          _compute_y(
              T1,
              W_with_dummy,
              beta,
              neural_network,
              stddev_outcome_noise,
          )
          - _compute_y(
              T0,
              W_with_dummy,
              beta,
              neural_network,
              stddev_outcome_noise,
          )
      )
  
      treatments = [("v" + str(i)) for i in range(0, num_treatments)]
      outcome = "y"
      common_causes = construct_col_names(
          "W",
          num_common_causes,
          num_discrete_common_causes,
          num_discrete_levels=4,
          one_hot_encode=one_hot_encode,
      )
      col_names = common_causes + treatments + [outcome]
      data = pd.DataFrame(data, columns=col_names)
      # Specifying the correct dtypes
      if treatment_is_binary:
          data = data.astype(
              {tname: "bool" for tname in treatments}, copy=False
          )
      elif treatment_is_category:
          data = data.astype(
              {tname: "category" for tname in treatments}, copy=False
          )
      if outcome_is_binary:
          data = data.astype({outcome: "bool"}, copy=False)
      if num_discrete_common_causes > 0 and not one_hot_encode:
          data = data.astype(
              {
                  wname: "int64"
                  for wname in common_causes[num_cont_common_causes:]
              },
              copy=False,
          )
          data = data.astype(
              {
                  wname: "category"
                  for wname in common_causes[num_cont_common_causes:]
              },
              copy=False,
          )
      dot_graph = create_dot_graph(
          treatments, outcome, common_causes, instruments=[]
      )
      # Now writing the gml graph
      gml_graph = create_gml_graph(
          treatments, outcome, common_causes, instruments=[]
      )
      ret_dict = {
          "df": data,
          "treatment_name": treatments,
          "outcome_name": outcome,
          "common_causes_names": common_causes,
          "dot_graph": dot_graph,
          "gml_graph": gml_graph,
          "ate": ate,
      }
      return ret_dict
  
imports: import math
  import string
  import networkx
  import numpy
  import pandas
  import scipy.stats
  from numpy.random import choice
  from scipy.stats import bernoulli, halfnorm, poisson, uniform
  from sklearn.neural_network import MLPRegressor
  from dowhy.utils.graph_operations import add_edge, del_edge, get_random_node_pair, get_simple_ordered_tree
label_desc: Use image node_id 15 for calling a global function with example usage: partially_linear_dataset(beta, num_common_causes, num_unobserved_common_causes, strength_unobserved_confounding, num_samples, num_treatments, treatment_is_binary, treatment_is_category, outcome_is_binary, stochastic_discretization, num_discrete_common_causes, stddev_treatment_noise, stddev_outcome_noise, one_hot_encode, training_sample_size, random_state) and returns: ret_dict, y