original_docstring: Train a neural network to solve an ODE.
  
      :param ode_system: The ODE system to solve. If the ODE system consists of equations :math:`F_i(x_1, x_2, ..., x_n, t) = 0` where :math:`x_i` is the dependent i-th variable and :math:`t` is the independent variable,
          then `ode_system` should be a function that maps :math:`(x_1, x_2, ..., x_n, t)` to a list where the i-th entry is :math:`F_i(x_1, x_2, ..., x_n, t)`.
      :type ode_system: function
      :param conditions: The initial/boundary conditions. The ith entry of the conditions is the condition that :math:`x_i` should satisfy.
      :type conditions: list[`neurodiffeq.ode.IVP` or `neurodiffeq.ode.DirichletBVP` or `neurodiffeq.ode.NoCondition`]
      :param t_min: The lower bound of the domain (t) on which the ODE is solved, only needed when train_generator or valid_generator are not specified, defaults to None
      :type t_min: float
      :param t_max: The upper bound of the domain (t) on which the ODE is solved, only needed when train_generator or valid_generator are not specified, defaults to None
      :type t_max: float
      :param single_net: The single neural network used to approximate the solution. Only one of `single_net` and `nets` should be specified, defaults to None
      :param single_net: `torch.nn.Module`, optional
      :param nets: The neural networks used to approximate the solution, defaults to None.
      :type nets: list[`torch.nn.Module`], optional
      :param train_generator: The example generator to generate 1-D training points, default to None.
      :type train_generator: `neurodiffeq.ode.ExampleGenerator`, optional
      :param shuffle: Whether to shuffle the training examples every epoch, defaults to True.
      :type shuffle: bool, optional
      :param valid_generator: The example generator to generate 1-D validation points, default to None.
      :type valid_generator: `neurodiffeq.ode.ExampleGenerator`, optional
      :param optimizer: The optimization method to use for training, defaults to None.
      :type optimizer: `torch.optim.Optimizer`, optional
      :param criterion: The loss function to use for training, defaults to None and sum of square of the output of `ode_system` will be used.
      :type criterion: `torch.nn.modules.loss._Loss`, optional
      :param additional_loss_term: Extra terms to add to the loss function besides the part specified by `criterion`. The input of `additional_loss_term` should be the same as `ode_system`
      :type additional_loss_term: function
      :param batch_size: The size of the mini-batch to use, defaults to 16.
      :type batch_size: int, optional
      :param max_epochs: The maximum number of epochs to train, defaults to 1000.
      :type max_epochs: int, optional
      :param monitor: The monitor to check the status of nerual network during training, defaults to None.
      :type monitor: `neurodiffeq.ode.Monitor`, optional
      :param return_internal: Whether to return the nets, conditions, training generator, validation generator, optimizer and loss function, defaults to False.
      :type return_internal: bool, optional
      :param return_best: Whether to return the nets that achieved the lowest validation loss, defaults to False.
      :type return_best: bool, optional
      :return: The solution of the ODE. The history of training loss and validation loss.
          Optionally, the nets, conditions, training generator, validation generator, optimizer and loss function.
      :rtype: tuple[`neurodiffeq.ode.Solution`, dict]; or tuple[`neurodiffeq.ode.Solution`, dict, dict]
      
docstring: Train a neural network to solve an ODE.
docstring_tokens: Train a neural network to solve an ODE .
code: def solve_system(
          ode_system, conditions, t_min, t_max,
          single_net=None, nets=None, train_generator=None, shuffle=True, valid_generator=None,
          optimizer=None, criterion=None, additional_loss_term=None, batch_size=16,
          max_epochs=1000,
          monitor=None, return_internal=False,
          return_best=False,
  ):
      """Train a neural network to solve an ODE.
  
      :param ode_system: The ODE system to solve. If the ODE system consists of equations :math:`F_i(x_1, x_2, ..., x_n, t) = 0` where :math:`x_i` is the dependent i-th variable and :math:`t` is the independent variable,
          then `ode_system` should be a function that maps :math:`(x_1, x_2, ..., x_n, t)` to a list where the i-th entry is :math:`F_i(x_1, x_2, ..., x_n, t)`.
      :type ode_system: function
      :param conditions: The initial/boundary conditions. The ith entry of the conditions is the condition that :math:`x_i` should satisfy.
      :type conditions: list[`neurodiffeq.ode.IVP` or `neurodiffeq.ode.DirichletBVP` or `neurodiffeq.ode.NoCondition`]
      :param t_min: The lower bound of the domain (t) on which the ODE is solved, only needed when train_generator or valid_generator are not specified, defaults to None
      :type t_min: float
      :param t_max: The upper bound of the domain (t) on which the ODE is solved, only needed when train_generator or valid_generator are not specified, defaults to None
      :type t_max: float
      :param single_net: The single neural network used to approximate the solution. Only one of `single_net` and `nets` should be specified, defaults to None
      :param single_net: `torch.nn.Module`, optional
      :param nets: The neural networks used to approximate the solution, defaults to None.
      :type nets: list[`torch.nn.Module`], optional
      :param train_generator: The example generator to generate 1-D training points, default to None.
      :type train_generator: `neurodiffeq.ode.ExampleGenerator`, optional
      :param shuffle: Whether to shuffle the training examples every epoch, defaults to True.
      :type shuffle: bool, optional
      :param valid_generator: The example generator to generate 1-D validation points, default to None.
      :type valid_generator: `neurodiffeq.ode.ExampleGenerator`, optional
      :param optimizer: The optimization method to use for training, defaults to None.
      :type optimizer: `torch.optim.Optimizer`, optional
      :param criterion: The loss function to use for training, defaults to None and sum of square of the output of `ode_system` will be used.
      :type criterion: `torch.nn.modules.loss._Loss`, optional
      :param additional_loss_term: Extra terms to add to the loss function besides the part specified by `criterion`. The input of `additional_loss_term` should be the same as `ode_system`
      :type additional_loss_term: function
      :param batch_size: The size of the mini-batch to use, defaults to 16.
      :type batch_size: int, optional
      :param max_epochs: The maximum number of epochs to train, defaults to 1000.
      :type max_epochs: int, optional
      :param monitor: The monitor to check the status of nerual network during training, defaults to None.
      :type monitor: `neurodiffeq.ode.Monitor`, optional
      :param return_internal: Whether to return the nets, conditions, training generator, validation generator, optimizer and loss function, defaults to False.
      :type return_internal: bool, optional
      :param return_best: Whether to return the nets that achieved the lowest validation loss, defaults to False.
      :type return_best: bool, optional
      :return: The solution of the ODE. The history of training loss and validation loss.
          Optionally, the nets, conditions, training generator, validation generator, optimizer and loss function.
      :rtype: tuple[`neurodiffeq.ode.Solution`, dict]; or tuple[`neurodiffeq.ode.Solution`, dict, dict]
      """
  
      ########################################### subroutines ###########################################
      def train(train_generator, net, nets, ode_system, conditions, criterion, additional_loss_term, shuffle, optimizer):
          train_examples_t = train_generator.get_examples()
          train_examples_t = train_examples_t.reshape((-1, 1))
          n_examples_train = train_generator.size
          idx = np.random.permutation(n_examples_train) if shuffle else np.arange(n_examples_train)
  
          train_loss_epoch = 0.0
          batch_start, batch_end = 0, batch_size
          while batch_start < n_examples_train:
              if batch_end > n_examples_train:
                  batch_end = n_examples_train
              batch_idx = idx[batch_start:batch_end]
              ts = train_examples_t[batch_idx]
  
              train_loss_batch = calculate_loss(ts, net, nets, ode_system, conditions, criterion, additional_loss_term)
              train_loss_epoch += train_loss_batch.item() * (batch_end - batch_start) / n_examples_train
  
              optimizer.zero_grad()
              train_loss_batch.backward()
              optimizer.step()
  
              batch_start += batch_size
              batch_end += batch_size
  
          return train_loss_epoch
  
      def valid(valid_generator, net, nets, ode_system, conditions, criterion, additional_loss_term):
          valid_examples_t = valid_generator.get_examples()
          ts = valid_examples_t.reshape((-1, 1))
          valid_loss_epoch = calculate_loss(ts, net, nets, ode_system, conditions, criterion, additional_loss_term)
          valid_loss_epoch = valid_loss_epoch.item()
          return valid_loss_epoch
  
      def calculate_loss(ts, net, nets, ode_system, conditions, criterion, additional_loss_term):
          us = _trial_solution(net, nets, ts, conditions)
          Futs = ode_system(*us, ts)
          loss = sum(
              criterion(Fut, torch.zeros_like(ts))
              for Fut in Futs
          )
          if additional_loss_term is not None:
              loss += additional_loss_term(*us, ts)
          return loss
  
      ###################################################################################################
  
      if single_net and nets:
          raise RuntimeError('Only one of net and nets should be specified')
      # defaults to use a single neural network
      if (not single_net) and (not nets):
          single_net = FCNN(n_input_units=1, n_output_units=len(conditions), n_hidden_units=32, n_hidden_layers=1,
                     actv=nn.Tanh)
      if single_net:
          # mark the Conditions so that we know which condition correspond to which output unit
          for ith, con in enumerate(conditions):
              con.set_impose_on(ith)
      if not train_generator:
          if (t_min is None) or (t_max is None):
              raise RuntimeError('Please specify t_min and t_max when train_generator is not specified')
          train_generator = ExampleGenerator(32, t_min, t_max, method='equally-spaced-noisy')
      if not valid_generator:
          if (t_min is None) or (t_max is None):
              raise RuntimeError('Please specify t_min and t_max when train_generator is not specified')
          valid_generator = ExampleGenerator(32, t_min, t_max, method='equally-spaced')
      if (not optimizer) and single_net:  # using a single net
          optimizer = optim.Adam(single_net.parameters(), lr=0.001)
      if (not optimizer) and nets:  # using multiple nets
          all_parameters = []
          for net in nets:
              all_parameters += list(net.parameters())
          optimizer = optim.Adam(all_parameters, lr=0.001)
      if not criterion:
          criterion = nn.MSELoss()
  
      loss_history = {'train': [], 'valid': []}
  
      if return_best:
          valid_loss_epoch_min = np.inf
          solution_min = None
  
      for epoch in range(max_epochs):
          train_loss_epoch = train(train_generator, single_net, nets, ode_system, conditions, criterion, additional_loss_term, shuffle,
                                   optimizer)
          loss_history['train'].append(train_loss_epoch)
  
          valid_loss_epoch = valid(valid_generator, single_net, nets, ode_system, conditions, criterion, additional_loss_term,)
          loss_history['valid'].append(valid_loss_epoch)
  
          if monitor and epoch % monitor.check_every == 0:
              monitor.check(single_net, nets, conditions, loss_history)
  
          if return_best and valid_loss_epoch < valid_loss_epoch_min:
              valid_loss_epoch_min = valid_loss_epoch
              solution_min = Solution(single_net, nets, conditions)
  
      if return_best:
          solution = solution_min
      else:
          solution = Solution(single_net, nets, conditions)
  
      if return_internal:
          internal = {
              'single_net': single_net,
              'nets': nets,
              'conditions': conditions,
              'train_generator': train_generator,
              'valid_generator': valid_generator,
              'optimizer': optimizer,
              'criterion': criterion
          }
          return solution, loss_history, internal
      else:
          return solution, loss_history
code_tokens: def solve_system ( ode_system , conditions , t_min , t_max , single_net = None , nets = None , train_generator = None , shuffle = True , valid_generator = None , optimizer = None , criterion = None , additional_loss_term = None , batch_size = 16 , max_epochs = 1000 , monitor = None , return_internal = False , return_best = False , ) : def train ( train_generator , net , nets , ode_system , conditions , criterion , additional_loss_term , shuffle , optimizer ) : train_examples_t = train_generator . get_examples ( ) train_examples_t = train_examples_t . reshape ( ( - 1 , 1 ) ) n_examples_train = train_generator . size idx = np . random . permutation ( n_examples_train ) if shuffle else np . arange ( n_examples_train ) train_loss_epoch = 0.0 batch_start , batch_end = 0 , batch_size while batch_start < n_examples_train : if batch_end > n_examples_train : batch_end = n_examples_train batch_idx = idx [ batch_start : batch_end ] ts = train_examples_t [ batch_idx ] train_loss_batch = calculate_loss ( ts , net , nets , ode_system , conditions , criterion , additional_loss_term ) train_loss_epoch += train_loss_batch . item ( ) * ( batch_end - batch_start ) / n_examples_train optimizer . zero_grad ( ) train_loss_batch . backward ( ) optimizer . step ( ) batch_start += batch_size batch_end += batch_size return train_loss_epoch def valid ( valid_generator , net , nets , ode_system , conditions , criterion , additional_loss_term ) : valid_examples_t = valid_generator . get_examples ( ) ts = valid_examples_t . reshape ( ( - 1 , 1 ) ) valid_loss_epoch = calculate_loss ( ts , net , nets , ode_system , conditions , criterion , additional_loss_term ) valid_loss_epoch = valid_loss_epoch . item ( ) return valid_loss_epoch def calculate_loss ( ts , net , nets , ode_system , conditions , criterion , additional_loss_term ) : us = _trial_solution ( net , nets , ts , conditions ) Futs = ode_system ( * us , ts ) loss = sum ( criterion ( Fut , torch . zeros_like ( ts ) ) for Fut in Futs ) if additional_loss_term is not None : loss += additional_loss_term ( * us , ts ) return loss if single_net and nets : raise RuntimeError ( 'Only one of net and nets should be specified' ) if ( not single_net ) and ( not nets ) : single_net = FCNN ( n_input_units = 1 , n_output_units = len ( conditions ) , n_hidden_units = 32 , n_hidden_layers = 1 , actv = nn . Tanh ) if single_net : for ith , con in enumerate ( conditions ) : con . set_impose_on ( ith ) if not train_generator : if ( t_min is None ) or ( t_max is None ) : raise RuntimeError ( 'Please specify t_min and t_max when train_generator is not specified' ) train_generator = ExampleGenerator ( 32 , t_min , t_max , method = 'equally-spaced-noisy' ) if not valid_generator : if ( t_min is None ) or ( t_max is None ) : raise RuntimeError ( 'Please specify t_min and t_max when train_generator is not specified' ) valid_generator = ExampleGenerator ( 32 , t_min , t_max , method = 'equally-spaced' ) if ( not optimizer ) and single_net : optimizer = optim . Adam ( single_net . parameters ( ) , lr = 0.001 ) if ( not optimizer ) and nets : all_parameters = [ ] for net in nets : all_parameters += list ( net . parameters ( ) ) optimizer = optim . Adam ( all_parameters , lr = 0.001 ) if not criterion : criterion = nn . MSELoss ( ) loss_history = { 'train' : [ ] , 'valid' : [ ] } if return_best : valid_loss_epoch_min = np . inf solution_min = None for epoch in range ( max_epochs ) : train_loss_epoch = train ( train_generator , single_net , nets , ode_system , conditions , criterion , additional_loss_term , shuffle , optimizer ) loss_history [ 'train' ] . append ( train_loss_epoch ) valid_loss_epoch = valid ( valid_generator , single_net , nets , ode_system , conditions , criterion , additional_loss_term , ) loss_history [ 'valid' ] . append ( valid_loss_epoch ) if monitor and epoch % monitor . check_every == 0 : monitor . check ( single_net , nets , conditions , loss_history ) if return_best and valid_loss_epoch < valid_loss_epoch_min : valid_loss_epoch_min = valid_loss_epoch solution_min = Solution ( single_net , nets , conditions ) if return_best : solution = solution_min else : solution = Solution ( single_net , nets , conditions ) if return_internal : internal = { 'single_net' : single_net , 'nets' : nets , 'conditions' : conditions , 'train_generator' : train_generator , 'valid_generator' : valid_generator , 'optimizer' : optimizer , 'criterion' : criterion } return solution , loss_history , internal else : return solution , loss_history
short_docstring: Train a neural network to solve an ODE.
short_docstring_tokens: Train a neural network to solve an ODE .
comment: """Train a neural network to solve an ODE.
  
      :param ode_system: The ODE system to solve. If the ODE system consists of equations :math:`F_i(x_1, x_2, ..., x_n, t) = 0` where :math:`x_i` is the dependent i-th variable and :math:`t` is the independent variable,
          then `ode_system` should be a function that maps :math:`(x_1, x_2, ..., x_n, t)` to a list where the i-th entry is :math:`F_i(x_1, x_2, ..., x_n, t)`.
      :type ode_system: function
      :param conditions: The initial/boundary conditions. The ith entry of the conditions is the condition that :math:`x_i` should satisfy.
      :type conditions: list[`neurodiffeq.ode.IVP` or `neurodiffeq.ode.DirichletBVP` or `neurodiffeq.ode.NoCondition`]
      :param t_min: The lower bound of the domain (t) on which the ODE is solved, only needed when train_generator or valid_generator are not specified, defaults to None
      :type t_min: float
      :param t_max: The upper bound of the domain (t) on which the ODE is solved, only needed when train_generator or valid_generator are not specified, defaults to None
      :type t_max: float
      :param single_net: The single neural network used to approximate the solution. Only one of `single_net` and `nets` should be specified, defaults to None
      :param single_net: `torch.nn.Module`, optional
      :param nets: The neural networks used to approximate the solution, defaults to None.
      :type nets: list[`torch.nn.Module`], optional
      :param train_generator: The example generator to generate 1-D training points, default to None.
      :type train_generator: `neurodiffeq.ode.ExampleGenerator`, optional
      :param shuffle: Whether to shuffle the training examples every epoch, defaults to True.
      :type shuffle: bool, optional
      :param valid_generator: The example generator to generate 1-D validation points, default to None.
      :type valid_generator: `neurodiffeq.ode.ExampleGenerator`, optional
      :param optimizer: The optimization method to use for training, defaults to None.
      :type optimizer: `torch.optim.Optimizer`, optional
      :param criterion: The loss function to use for training, defaults to None and sum of square of the output of `ode_system` will be used.
      :type criterion: `torch.nn.modules.loss._Loss`, optional
      :param additional_loss_term: Extra terms to add to the loss function besides the part specified by `criterion`. The input of `additional_loss_term` should be the same as `ode_system`
      :type additional_loss_term: function
      :param batch_size: The size of the mini-batch to use, defaults to 16.
      :type batch_size: int, optional
      :param max_epochs: The maximum number of epochs to train, defaults to 1000.
      :type max_epochs: int, optional
      :param monitor: The monitor to check the status of nerual network during training, defaults to None.
      :type monitor: `neurodiffeq.ode.Monitor`, optional
      :param return_internal: Whether to return the nets, conditions, training generator, validation generator, optimizer and loss function, defaults to False.
      :type return_internal: bool, optional
      :param return_best: Whether to return the nets that achieved the lowest validation loss, defaults to False.
      :type return_best: bool, optional
      :return: The solution of the ODE. The history of training loss and validation loss.
          Optionally, the nets, conditions, training generator, validation generator, optimizer and loss function.
      :rtype: tuple[`neurodiffeq.ode.Solution`, dict]; or tuple[`neurodiffeq.ode.Solution`, dict, dict]
      """$SEP$########################################### subroutines ###########################################$SEP$###################################################################################################$SEP$# defaults to use a single neural network$SEP$# mark the Conditions so that we know which condition correspond to which output unit$SEP$# using a single net$SEP$# using multiple nets