===================================================
Demo: Using the Block Solver with Nested Splits
===================================================

In :doc:`demo_block_solver_nse`, we demonstrated how to use the block solver
to solve the steady Navier-Stokes equations. In that example, we only needed a single split 
to define the block structure of the system. In this demo, we will show how to use nested splits
to define a more complex block structure.

Problem Setup
-----------------
First, we define a simple scalar Galerkin Probject Physics Object to test the block solver with nested splits.

.. code-block:: python

    class GalerkinProjection(PhysicsProblem):
        """
        GalerkinProjection field_value = b
        """
        def set_projection_value(self, projection_value):
            self.set_external_function('b', projection_value)

        def flux(self):
            ''''''

        def get_residual(self):
            ''''''

        def set_weak_form(self):
            b = self.external_function('b')
            u = self.get_solution_function()
            w = self.get_test_function()
            self.weak_form = ufl.inner(u-b, w)*self.dx

Then we define a helper function to create three GalerkinProjection that we will couple together.

.. code-block:: python

    def build_GP(tag, mesh, val):
        GP = GalerkinProjection(mesh, tag=tag)
        GP.set_element('CG', 1)
        GP.set_projection_value(dolfinx.fem.Constant(mesh.msh, dolfinx.default_scalar_type(val)))
        return GP 

Next, we create a mesh and three GalerkinProjection objects using the helper function.

.. code-block:: python
    
    # Create a mesh
    mesh = LineMesh(0, 1, 1/5)
    
    # Define three GP physics with different projection values
    GP1 = build_GP('A', mesh, val=0)
    GP2 = build_GP('B', mesh, val=1)
    GP3 = build_GP('C', mesh, val=2)
    GPs = [GP1, GP2, GP3]

We then monolithically couple the three GalerkinProjection objects using the MultiphysicsProblem class.

.. code-block:: python

    # Define a multiphysics problem
    phs = MultiphysicsProblem(*GPs)
    phs.set_element()
    phs.build_function_space()
    phs.set_weak_form()

    # Define the problem 
    problem = NonLinearProblem(phs)

Next, we define the block structure of the system using nested splits. The first split GP2 from GP1 and GP3, and then we split GP1 and GP3 individually.
Here, we use the default ksp objects for each split. For an example of how to customize the ksp objects, see :doc:`demo_block_solver_nse`.

.. code-block:: python

    # Split the problem into two blocks: (A,C) and (B)
    split0 = {'fields': (('A','C'), 'B'), 
                'composite_type': 'schur', 
                'schur_fact_type': 'full', 
                'schur_pre_type': 'a11'}

    # Second split: (A,C) into (A) and (C)
    split1 = {'fields': ('A','C'), 
                'composite_type': 'schur', 
                'schur_fact_type': 'full', 
                'schur_pre_type': 'a11'}

    # Create a list of splits to define a tree structure
    splits = [split0, split1]

    # Define the tree and solver
    tree = BlockSplitTree(phs, splits=splits)
    solver = BlockNonLinearSolver(tree, MPI.COMM_WORLD, problem)

    # Solve the problem
    solver.solve()

Finally, we verify that the solution is correct.

.. code-block:: python
    
    # Split the solution into its components 
    A, B, C = phs.get_solution_function().split()

    # Get the dofs for each component
    dofs_A = phs.get_function_space().sub(0).dofmap.list.flatten()
    dofs_B = phs.get_function_space().sub(1).dofmap.list.flatten()
    dofs_C = phs.get_function_space().sub(2).dofmap.list.flatten()

    # Check that the solution is correct
    print("Mesh vertices:", mesh.msh.geometry.x.ravel())
    print("A values at vertices:", A.x.array[dofs_A])
    print("B values at vertices:", B.x.array[dofs_B])
    print("C values at vertices:", C.x.array[dofs_C])

Full Script
-----------------

.. code-block:: python

    import dolfinx
    import ufl

    from flatiron_tk.mesh import LineMesh
    from flatiron_tk.physics import MultiphysicsProblem
    from flatiron_tk.physics import PhysicsProblem
    from flatiron_tk.solver import NonLinearProblem
    from flatiron_tk.solver import BlockSplitTree
    from flatiron_tk.solver import BlockNonLinearSolver
    from mpi4py import MPI

    # Build GP physics
    class GalerkinProjection(PhysicsProblem):
        """
        GalerkinProjection field_value = b
        """
        def set_projection_value(self, projection_value):
            self.set_external_function('b', projection_value)

        def flux(self):
            ''''''

        def get_residual(self):
            ''''''

        def set_weak_form(self):
            b = self.external_function('b')
            u = self.get_solution_function()
            w = self.get_test_function()
            self.weak_form = ufl.inner(u-b, w)*self.dx

    # Helper function to build GP physics
    def build_GP(tag, mesh, val):
        GP = GalerkinProjection(mesh, tag=tag)
        GP.set_element('CG', 1)
        GP.set_projection_value(dolfinx.fem.Constant(mesh.msh, dolfinx.default_scalar_type(val)))
        return GP 

    # Create a mesh
    mesh = LineMesh(0, 1, 1/5)

    # Define three GP physics with different projection values
    GP1 = build_GP('A', mesh, val=0)
    GP2 = build_GP('B', mesh, val=1)
    GP3 = build_GP('C', mesh, val=2)
    GPs = [GP1, GP2, GP3]

    # Define a multiphysics problem
    phs = MultiphysicsProblem(*GPs)
    phs.set_element()
    phs.build_function_space()
    phs.set_weak_form()

    # Define the problem 
    problem = NonLinearProblem(phs)

    # Split the problem into two blocks: (A,C) and (B)
    split0 = {'fields': (('A','C'), 'B'), 
                'composite_type': 'schur', 
                'schur_fact_type': 'full', 
                'schur_pre_type': 'a11'}

    # Second split: (A,C) into (A) and (C)
    split1 = {'fields': ('A','C'), 
                'composite_type': 'schur', 
                'schur_fact_type': 'full', 
                'schur_pre_type': 'a11'}

    # Create a list of splits to define a tree structure
    splits = [split0, split1]

    # Define the tree and solver
    tree = BlockSplitTree(phs, splits=splits)
    solver = BlockNonLinearSolver(tree, MPI.COMM_WORLD, problem)

    # Solve the problem
    solver.solve()

    # Split the solution into its components 
    A, B, C = phs.get_solution_function().split()

    # Get the dofs for each component
    dofs_A = phs.get_function_space().sub(0).dofmap.list.flatten()
    dofs_B = phs.get_function_space().sub(1).dofmap.list.flatten()
    dofs_C = phs.get_function_space().sub(2).dofmap.list.flatten()

    # Check that the solution is correct
    print("Mesh vertices:", mesh.msh.geometry.x.ravel())
    print("A values at vertices:", A.x.array[dofs_A])
    print("B values at vertices:", B.x.array[dofs_B])
    print("C values at vertices:", C.x.array[dofs_C])