Repository: stfc/PSyclone

Issue:
psyad assignment transformation does not work with array ranges
I'm sure it used to work but assignments with array ranges now causes psyad to complain that the assignment is not in tangent-linear form. I have had to change the example to avoid this problem.

Discussion:

Patch:
diff --git a/changelog b/changelog
index 1f926e56a..93ad03460 100644
--- a/changelog
+++ b/changelog
@@ -67,6 +67,9 @@
 	23) PR #1228 towards #753. The GOcean API now uses PSyIR for the Algorithm
 	layer.
 
+	24) PR #1526 for #1430. Fixes and improves support for array
+	ranges in psyad.
+
 release 2.1.0 3rd November 2021
 
 	1) PR #1183 for #1075. Improves (tightens) the checking performed
diff --git a/examples/psyad/eg2/Makefile b/examples/psyad/eg2/Makefile
index 920bf4327..b97758cc2 100644
--- a/examples/psyad/eg2/Makefile
+++ b/examples/psyad/eg2/Makefile
@@ -31,7 +31,7 @@
 # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 # POSSIBILITY OF SUCH DAMAGE.
 # ------------------------------------------------------------------------------
-# Author: R. W. Ford, STFC Daresbury Laboratory
+# Authors: R. W. Ford and A. R. Porter, STFC Daresbury Laboratory
 
 include ../../common.mk
 
@@ -42,7 +42,7 @@ GENERATED_FILES +=
 # second, if supported, would write both it and the corresponding test
 # harness to separate files.
 transform:
-	$(PSYAD) tl_hydrostatic_kernel_mod.F90 -a r_u exner theta moist_dyn_gas moist_dyn_tot moist_dyn_fac grad_term theta_v_e theta_v_at_quad exner_e exner_at_quad
+	$(PSYAD) tl_hydrostatic_kernel_mod.F90 -a r_u exner theta moist_dyn_gas moist_dyn_tot moist_dyn_fac grad_term theta_v_e theta_v_at_quad grad_theta_v_at_quad exner_e exner_at_quad
 	@echo "Harness generation not yet supported"
 
 # It is not yet possible to compile the adjoint code into an object
diff --git a/examples/psyad/eg2/tl_hydrostatic_kernel_mod.F90 b/examples/psyad/eg2/tl_hydrostatic_kernel_mod.F90
index 6c08575ed..c58e04099 100644
--- a/examples/psyad/eg2/tl_hydrostatic_kernel_mod.F90
+++ b/examples/psyad/eg2/tl_hydrostatic_kernel_mod.F90
@@ -248,19 +248,16 @@ subroutine tl_hydrostatic_code(nlayers,          &
           exner_at_quad  = exner_at_quad + exner_e(df)*w3_basis(1,df,qp1,qp2)
         end do
         theta_v_at_quad = 0.0_r_def
-        !*** Issue #1430 array notation for active variables is not working
-        !*** grad_theta_v_at_quad(:) = 0.0_r_def
+        grad_theta_v_at_quad(:) = 0.0_r_def
         do df = 1, ndf_wt
           theta_v_at_quad   = theta_v_at_quad                                 &
                             + theta_v_e(df)*wt_basis(1,df,qp1,qp2)
-          !*** Issue #1430 array notation for active variables is not working
-          !*** grad_theta_v_at_quad(:) = grad_theta_v_at_quad(:)                   &
-          !***                         + theta_v_e(df)*wt_diff_basis(:,df,qp1,qp2)
+          grad_theta_v_at_quad(:) = grad_theta_v_at_quad(:)                   &
+                                   + theta_v_e(df)*wt_diff_basis(:,df,qp1,qp2)
         end do
         ! Calculation
         do df = 1, ndf_w2
-          !*** Issue #1430 array notation for active variables is not working
-          !*** v  = w2_basis(:,df,qp1,qp2)
+          v  = w2_basis(:,df,qp1,qp2)
           dv = w2_diff_basis(1,df,qp1,qp2)
           ! Pressure gradient term
           !*** Issue #1490 dot_product is not supported
diff --git a/src/psyclone/psyad/transformations/assignment_trans.py b/src/psyclone/psyad/transformations/assignment_trans.py
index 01100b8bc..6ae71f96d 100644
--- a/src/psyclone/psyad/transformations/assignment_trans.py
+++ b/src/psyclone/psyad/transformations/assignment_trans.py
@@ -1,6 +1,6 @@
 # BSD 3-Clause License
 #
-# Copyright (c) 2021, Science and Technology Facilities Council.
+# Copyright (c) 2021-2022, Science and Technology Facilities Council.
 # All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
@@ -42,6 +42,7 @@ from __future__ import absolute_import
 from psyclone.core import SymbolicMaths
 from psyclone.psyir.nodes import BinaryOperation, Assignment, Reference, \
     Literal, UnaryOperation
+from psyclone.psyir.nodes.array_mixin import ArrayMixin
 from psyclone.psyir.symbols import REAL_TYPE
 from psyclone.psyir.transformations import TransformationError
 
@@ -89,15 +90,28 @@ class AssignmentTrans(AdjointTransformation):
             active_var = None
             new_rhs_term = rhs_term.copy()
             for ref in new_rhs_term.walk(Reference):
-                if ref.symbol in self._active_variables:
-                    active_var = ref
+                if ref.symbol not in self._active_variables:
+                    continue
+                active_var = ref
+                # Identify whether this reference on the RHS matches the
+                # one on the LHS - if so we have an increment.
+                if node.is_array_range and isinstance(ref, ArrayMixin):
+                    # TODO #1537 - we can't just do `sym_maths.equal` if we
+                    # have an array range because the SymbolicMaths class does
+                    # not currently support them.
+                    # Since we have already checked (in validate) that any
+                    # references to the same symbol on the RHS have the same
+                    # range, this is an increment if the symbols match.
+                    if node.lhs.symbol is ref.symbol:
+                        increment = True
+                else:
                     if sym_maths.equal(ref, node.lhs):
                         increment = True
-                    if ref.parent:
-                        ref.replace_with(node.lhs.copy())
-                    else:
-                        new_rhs_term = node.lhs.copy()
-                    break
+                if ref.parent:
+                    ref.replace_with(node.lhs.copy())
+                else:
+                    new_rhs_term = node.lhs.copy()
+                break
 
             # Work out whether the binary operation for this term is a
             # '+' or a '-' and return it in 'rhs_operator'.
@@ -159,6 +173,59 @@ class AssignmentTrans(AdjointTransformation):
         # Remove the original node
         node.detach()
 
+    def _array_ranges_match(self, assign, active_variable):
+        '''
+        If the supplied assignment is to an array range and the supplied
+        active variable is the entity being assigned to then this routine
+        checks that the array ranges of the LHS and the supplied reference
+        match. If they do not then an exception is raised.
+
+        :param assign: the assignment that we are checking.
+        :type assign: :py:class:`psyclone.psyir.nodes.Assignment`
+        :param active_variable: an active variable that appears on the \
+            LHS and RHS of the supplied assignment.
+        :type active_variable: :py:class:`psyclone.psyir.nodes.Reference`
+
+        :raises TangentLinearError: if the supplied assignment is to a \
+            symbol with an array range but the same symbol occurs on the \
+            RHS without an array range.
+        :raises NotImplementedError: if the array ranges on the LHS and \
+            RHS of the assignment for the supplied variable do not match.
+
+        '''
+        # This check only needs to proceed if the assignment is to an array
+        # range and the supplied active variable is the one being assigned to.
+        if not (assign.is_array_range and active_variable.symbol is
+                assign.lhs.symbol):
+            return
+
+        if not isinstance(active_variable, ArrayMixin):
+            raise TangentLinearError(
+                f"Assignment is to an array range but found a "
+                f"reference to the LHS variable "
+                f"'{assign.lhs.symbol.name}' without array notation"
+                f" on the RHS: '{self._writer(assign)}'")
+
+        sym_maths = SymbolicMaths.get()
+
+        for pos, idx in enumerate(active_variable.indices):
+            lhs_idx = assign.lhs.indices[pos]
+            # TODO #1537. This is a workaround until the SymbolicMaths
+            # class supports the comparison of array ranges.
+            # pylint: disable=unidiomatic-typecheck
+            if not (type(idx) == type(lhs_idx) and
+                    sym_maths.equal(idx.start,
+                                    lhs_idx.start) and
+                    sym_maths.equal(idx.stop,
+                                    lhs_idx.stop) and
+                    sym_maths.equal(idx.step,
+                                    lhs_idx.step)):
+                raise NotImplementedError(
+                    f"Different sections of the same active array "
+                    f"'{assign.lhs.symbol.name}' are "
+                    f"accessed on the LHS and RHS of an assignment: "
+                    f"'{self._writer(assign)}'. This is not supported.")
+
     def validate(self, node, options=None):
         '''Perform various checks to ensure that it is valid to apply the
         AssignmentTrans transformation to the supplied PSyIR Node.
@@ -179,10 +246,11 @@ class AssignmentTrans(AdjointTransformation):
             raise TransformationError(
                 "Node argument in assignment transformation should be a PSyIR "
                 "Assignment, but found '{0}'.".format(type(node).__name__))
+        assign = node
 
         # If there are no active variables then return
         assignment_active_var_names = [
-            var.name for var in node.walk(Reference)
+            var.name for var in assign.walk(Reference)
             if var.symbol in self._active_variables]
         if not assignment_active_var_names:
             # No active variables in this assigment so the assignment
@@ -190,18 +258,18 @@ class AssignmentTrans(AdjointTransformation):
             return
 
         # The lhs of the assignment node should be an active variable
-        if node.lhs.symbol not in self._active_variables:
+        if assign.lhs.symbol not in self._active_variables:
             # There are active vars on RHS but not on LHS
             raise TangentLinearError(
                 "Assignment node '{0}' has the following active variables on "
                 "its RHS '{1}' but its LHS '{2}' is not an active variable."
-                "".format(self._writer(node), assignment_active_var_names,
-                          node.lhs.name))
+                "".format(self._writer(assign), assignment_active_var_names,
+                          assign.lhs.name))
 
         # Split the RHS of the assignment into <expr> +- <expr> +- <expr>
         rhs_terms = self._split_nodes(
-            node.rhs, [BinaryOperation.Operator.ADD,
-                       BinaryOperation.Operator.SUB])
+            assign.rhs, [BinaryOperation.Operator.ADD,
+                         BinaryOperation.Operator.SUB])
 
         # Check for the special case where RHS=0.0. This is really a
         # representation of multiplying an active variable by zero but
@@ -215,27 +283,37 @@ class AssignmentTrans(AdjointTransformation):
         # A */ <expr> where A is an active variable.
         for rhs_term in rhs_terms:
 
-            active_vars = [
-                ref for ref in rhs_term.walk(Reference) if ref.symbol
-                in self._active_variables]
+            # When searching for references to an active variable we must
+            # take care to exclude those cases where they are present as
+            # arguments to the L/UBOUND intrinsics (as they will be when
+            # array notation is used).
+            active_vars = []
+            lu_bound_ops = [BinaryOperation.Operator.LBOUND,
+                            BinaryOperation.Operator.UBOUND]
+            for ref in rhs_term.walk(Reference):
+                if (ref.symbol in self._active_variables and
+                        not (isinstance(ref.parent, BinaryOperation) and
+                             ref.parent.operator in lu_bound_ops)):
+                    active_vars.append(ref)
 
             if not active_vars:
                 # This term must contain an active variable
                 raise TangentLinearError(
                     "Each non-zero term on the RHS of the assigment '{0}' "
                     "must have an active variable but '{1}' does not."
-                    "".format(self._writer(node), self._writer(rhs_term)))
+                    "".format(self._writer(assign), self._writer(rhs_term)))
 
             if len(active_vars) > 1:
                 # This term can only contain one active variable
                 raise TangentLinearError(
                     "Each term on the RHS of the assigment '{0}' must not "
                     "have more than one active variable but '{1}' has {2}."
-                    "".format(self._writer(node), self._writer(rhs_term),
+                    "".format(self._writer(assign), self._writer(rhs_term),
                               len(active_vars)))
 
-            if isinstance(rhs_term, Reference) and rhs_term.symbol \
-               in self._active_variables:
+            if (isinstance(rhs_term, Reference) and rhs_term.symbol
+                    in self._active_variables):
+                self._array_ranges_match(assign, rhs_term)
                 # This term consists of a single active variable (with
                 # a multiplier of unity) and is therefore valid.
                 continue
@@ -262,7 +340,7 @@ class AssignmentTrans(AdjointTransformation):
                     "Each term on the RHS of the assignment '{0}' must "
                     "be linear with respect to the active variable, but "
                     "found '{1}'.".format(
-                        self._writer(node), self._writer(rhs_term)))
+                        self._writer(assign), self._writer(rhs_term)))
 
             # The term must be a product of an active variable with an
             # inactive expression. Check that the active variable does
@@ -281,11 +359,16 @@ class AssignmentTrans(AdjointTransformation):
                         "In tangent-linear code an active variable cannot "
                         "appear as a denominator but '{0}' was found in "
                         "'{1}'.".format(
-                            self._writer(rhs_term), self._writer(node)))
+                            self._writer(rhs_term), self._writer(assign)))
                 # Continue up the PSyIR tree
                 candidate = parent
                 parent = candidate.parent
 
+            # If the LHS of the assignment is an array range then we only
+            # support accesses of the same variable on the RHS if they have
+            # the same range.
+            self._array_ranges_match(assign, active_variable)
+
     @staticmethod
     def _split_nodes(node, binary_operator_list):
         '''Utility to split an expression into a series of sub-expressions


Test patch:
diff --git a/src/psyclone/tests/psyad/transformations/test_assignment_trans.py b/src/psyclone/tests/psyad/transformations/test_assignment_trans.py
index 67a2b7b9c..87a52c33f 100644
--- a/src/psyclone/tests/psyad/transformations/test_assignment_trans.py
+++ b/src/psyclone/tests/psyad/transformations/test_assignment_trans.py
@@ -1,6 +1,6 @@
 # BSD 3-Clause License
 #
-# Copyright (c) 2021, Science and Technology Facilities Council.
+# Copyright (c) 2021-2022, Science and Technology Facilities Council.
 # All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
@@ -43,13 +43,20 @@ from psyclone.psyad.transformations import AssignmentTrans, TangentLinearError
 from psyclone.psyir.backend.fortran import FortranWriter
 from psyclone.psyir.frontend.fortran import FortranReader
 from psyclone.psyir.nodes import BinaryOperation, Reference, Assignment, \
-    Literal, UnaryOperation
+    Literal, UnaryOperation, ArrayReference, Range
 from psyclone.psyir.symbols import DataSymbol, REAL_TYPE, INTEGER_TYPE, \
-    ScalarType
+    ScalarType, ArrayType
 from psyclone.psyir.transformations import TransformationError
 from psyclone.tests.utilities import Compile
 
 
+@pytest.fixture(name="index_str", params=["1", "i", ":", "2:3"])
+def index_str_fixture(request):
+    ''' Fixture that supplies a number of different types of Fortran
+    array-index expressions. Returns the contents of params in turn. '''
+    return request.param
+
+
 def check_adjoint(tl_fortran, active_variable_names, expected_ad_fortran,
                   tmpdir):
     '''Utility routine that takes tangent linear fortran code as input in
@@ -71,10 +78,9 @@ def check_adjoint(tl_fortran, active_variable_names, expected_ad_fortran,
 
     '''
     # Add "subroutine / end subroutine" lines to the incoming code.
-    input_code = ("subroutine test()\n{0}end subroutine test\n"
-                  "".format(tl_fortran))
-    expected_output_code = ("subroutine test()\n{0}end subroutine test\n"
-                            "".format(expected_ad_fortran))
+    input_code = f"subroutine test()\n{tl_fortran}end subroutine test\n"
+    expected_output_code = (f"subroutine test()\n{expected_ad_fortran}end "
+                            f"subroutine test\n")
 
     # Translate the tangent linear code to PSyIR.
     reader = FortranReader()
@@ -115,10 +121,9 @@ def check_adjoint(tl_fortran, active_variable_names, expected_ad_fortran,
 def test_zero(tmpdir):
     '''Test the adjoint transformation with an assignment of the form A =
     0. This is the only valid value that an active variable can be set
-    to in a tangent linear code and represents multiplying an active
+    to in a tangent-linear code and represents multiplying an active
     variable by zero on the rhs. Scalars, directly addressed arrays,
-    indirectly addressed arrays and structure array accesses are
-    tested.
+    indirectly addressed arrays and array ranges are tested.
 
     A=0 -> A*=0
 
@@ -156,6 +161,28 @@ def test_zero(tmpdir):
         "  integer, dimension(n) :: b\n\n"
         "  a(b(n)) = 0.0\n\n")
     check_adjoint(tl_fortran, active_variables, ad_fortran, tmpdir)
+    # Array range
+    tl_fortran = (
+        "  integer, parameter :: n=10\n"
+        "  real :: a(n)\n"
+        "  a(:) = 0.0\n\n")
+    active_variables = ["a"]
+    ad_fortran = (
+        "  integer, parameter :: n = 10\n"
+        "  real, dimension(n) :: a\n\n"
+        "  a(:) = 0.0\n\n")
+    check_adjoint(tl_fortran, active_variables, ad_fortran, tmpdir)
+    # Array range with limits
+    tl_fortran = (
+        "  integer, parameter :: n=10\n"
+        "  real :: a(n)\n"
+        "  a(2:n-2) = 0.0\n\n")
+    active_variables = ["a"]
+    ad_fortran = (
+        "  integer, parameter :: n = 10\n"
+        "  real, dimension(n) :: a\n\n"
+        "  a(2:n - 2) = 0.0\n\n")
+    check_adjoint(tl_fortran, active_variables, ad_fortran, tmpdir)
 
 
 @pytest.mark.xfail(reason="issue #1347, type declaration is being written in "
@@ -190,8 +217,8 @@ def test_single_assign(tmpdir):
     '''Test the transformation works when there is one active variable on
     the rhs (B) and with the active variable on the lhs (A) only being
     written to, i.e. not also read on the rhs. Scalars, directly
-    addressed arrays, indirectly addressed arrays and structure array
-    accesses are tested.
+    addressed arrays, indirectly addressed arrays and array ranges are
+    tested.
 
     A=B -> B*=B*+A*; A*=0.0
 
@@ -234,6 +261,32 @@ def test_single_assign(tmpdir):
         "  b(lookup(n) + 1) = b(lookup(n) + 1) + a(lookup(2 * i))\n"
         "  a(lookup(2 * i)) = 0.0\n\n")
     check_adjoint(tl_fortran, active_variables, ad_fortran, tmpdir)
+    # Array range
+    tl_fortran = (
+        "  integer, parameter :: n=10\n"
+        "  real :: a(n),b(n)\n"
+        "  a(:) = b(:)\n")
+    active_variables = ["a", "b"]
+    ad_fortran = (
+        "  integer, parameter :: n = 10\n"
+        "  real, dimension(n) :: a\n"
+        "  real, dimension(n) :: b\n\n"
+        "  b(:) = b(:) + a(:)\n"
+        "  a(:) = 0.0\n\n")
+    check_adjoint(tl_fortran, active_variables, ad_fortran, tmpdir)
+    # Array range with limits
+    tl_fortran = (
+        "  integer, parameter :: n=10\n"
+        "  real :: a(n),b(n)\n"
+        "  a(2:n) = b(1:n-1)\n")
+    active_variables = ["a", "b"]
+    ad_fortran = (
+        "  integer, parameter :: n = 10\n"
+        "  real, dimension(n) :: a\n"
+        "  real, dimension(n) :: b\n\n"
+        "  b(1:n - 1) = b(1:n - 1) + a(2:n)\n"
+        "  a(2:n) = 0.0\n\n")
+    check_adjoint(tl_fortran, active_variables, ad_fortran, tmpdir)
 
 
 @pytest.mark.xfail(reason="issue #1347, type declaration is being written in "
@@ -317,7 +370,29 @@ def test_multi_add(tmpdir):
     check_adjoint(tl_fortran, active_variables, ad_fortran, tmpdir)
 
 
-def test_increment(tmpdir):
+def test_multi_add_array_elements(tmpdir):
+    '''
+    Same as test_multi_add except that 'A' and 'B' are actually just
+    different elements of the same array:
+
+    A=B+xC -> B*=B*+A*; C*=C*+xA*; A*=0.0
+
+    '''
+    tl_fortran = (
+        "  real a(10), b(10), n\n"
+        "  integer :: i,j\n"
+        "  a(i) = a(j) + 3*n*b(j)\n")
+    active_variables = ["a", "b"]
+    ad_fortran = (
+        "  real, dimension(10) :: a\n  real, dimension(10) :: b\n"
+        "  real :: n\n  integer :: i\n  integer :: j\n\n"
+        "  a(j) = a(j) + a(i)\n"
+        "  b(j) = b(j) + 3 * n * a(i)\n"
+        "  a(i) = 0.0\n\n")
+    check_adjoint(tl_fortran, active_variables, ad_fortran, tmpdir)
+
+
+def test_increment(tmpdir, index_str):
     '''Test the adjoint transformation with an assignment of the form
     A = A.
 
@@ -327,17 +402,19 @@ def test_increment(tmpdir):
 
     '''
     tl_fortran = (
-        "  integer, parameter :: n=2\n"
-        "  real a(n)\n"
-        "  a(n) = a(n)\n")
+        f"  integer, parameter :: n=3\n"
+        f"  integer :: i\n"
+        f"  real a(n)\n"
+        f"  a({index_str}) = a({index_str})\n")
     active_variables = ["a"]
     ad_fortran = (
-        "  integer, parameter :: n = 2\n"
+        "  integer, parameter :: n = 3\n"
+        "  integer :: i\n"
         "  real, dimension(n) :: a\n\n\n")
     check_adjoint(tl_fortran, active_variables, ad_fortran, tmpdir)
 
 
-def test_increment_mult(tmpdir):
+def test_increment_mult(tmpdir, index_str):
     '''Test the transformation works when the active variable that is
     written to on the lhs (A) is also read (and scaled (x)) on the rhs
     and there are no other active variables on the rhs. Also test
@@ -347,18 +424,20 @@ def test_increment_mult(tmpdir):
 
     '''
     tl_fortran = (
-        "  integer, parameter :: n=4\n"
-        "  real a(n)\n"
-        "  A(n) = 5*A(n)\n")
+        f"  integer, parameter :: n=4\n"
+        f"  integer :: i\n"
+        f"  real a(n)\n"
+        f"  A({index_str}) = 5*A({index_str})\n")
     active_variables = ["a"]
     ad_fortran = (
-        "  integer, parameter :: n = 4\n"
-        "  real, dimension(n) :: a\n\n"
-        "  a(n) = 5 * a(n)\n\n")
+        f"  integer, parameter :: n = 4\n"
+        f"  integer :: i\n"
+        f"  real, dimension(n) :: a\n\n"
+        f"  a({index_str}) = 5 * a({index_str})\n\n")
     check_adjoint(tl_fortran, active_variables, ad_fortran, tmpdir)
 
 
-def test_increment_add(tmpdir):
+def test_increment_add(tmpdir, index_str):
     '''Test the transformation works when the active variable that is
     written to on the lhs (A) is also read on the rhs and there is a
     another active variable (B) read on the rhs. Also test mixed-case
@@ -368,46 +447,50 @@ def test_increment_add(tmpdir):
 
     '''
     tl_fortran = (
-        "  real a(10), b(10)\n"
-        "  a(1) = A(1)+B(1)\n")
+        f"  real a(10), b(10)\n  integer :: i\n"
+        f"  a({index_str}) = A({index_str})+B({index_str})\n")
     active_variables = ["a", "b"]
     ad_fortran = (
-        "  real, dimension(10) :: a\n"
-        "  real, dimension(10) :: b\n\n"
-        "  b(1) = b(1) + a(1)\n\n")
+        f"  real, dimension(10) :: a\n"
+        f"  real, dimension(10) :: b\n"
+        f"  integer :: i\n\n"
+        f"  b({index_str}) = b({index_str}) + a({index_str})\n\n")
     check_adjoint(tl_fortran, active_variables, ad_fortran, tmpdir)
 
 
-def test_increment_add_reorder(tmpdir):
+def test_increment_add_reorder(tmpdir, index_str):
     '''Test the transformation works when the active variable that is
     written to on the lhs (A) is also read (and scaled (k)) on the rhs
     and there is another active variable (B) read on the
     rhs. Additionally, the active variable A is not the first term on
     the rhs (where we define terms as expressions separated by an
-    addition (or subtraction)).
+    addition (or subtraction)). Test for both individual array accesses
+    and array ranges.
 
     A=B+kA -> B*+=A*; A*=kA*
 
     '''
     tl_fortran = (
-        "  real a(10), b(10)\n"
-        "  integer k\n"
-        "  a(1) = b(1)+k*a(1)\n")
+        f"  real a(10), b(10)\n"
+        f"  integer i, k\n"
+        f"  a({index_str}) = b({index_str})+k*a({index_str})\n")
     active_variables = ["a", "b"]
     ad_fortran = (
-        "  real, dimension(10) :: a\n"
-        "  real, dimension(10) :: b\n"
-        "  integer :: k\n\n"
-        "  b(1) = b(1) + a(1)\n"
-        "  a(1) = k * a(1)\n\n")
+        f"  real, dimension(10) :: a\n"
+        f"  real, dimension(10) :: b\n"
+        f"  integer :: i\n"
+        f"  integer :: k\n\n"
+        f"  b({index_str}) = b({index_str}) + a({index_str})\n"
+        f"  a({index_str}) = k * a({index_str})\n\n".format(index_str))
     check_adjoint(tl_fortran, active_variables, ad_fortran, tmpdir)
 
 
-def test_increment_multi_add(tmpdir):
+def test_increment_multi_add(tmpdir, index_str):
     '''Test the transformation works when the active variable that is
     written to on the lhs (A) is also read (and scaled (w)) on the rhs
     and there are many other active variables on the rhs (B, C, D),
-    which are also scaled.
+    which are also scaled. Test for both individual array accesses
+    and array ranges.
 
     A=wA+xB+yC+zD -> D*=D*+zA*; C*=C*+yA*; B*=B*+xA*; A*=wA*
 
@@ -415,38 +498,44 @@ def test_increment_multi_add(tmpdir):
     tl_fortran = (
         "  real a(10), b(10), c(10), d(10)\n"
         "  real w(10), x, y(10), z\n"
-        "  a(1) = w(1)*a(1)+x*b(1)+y(1)*c(1)+d(1)*z\n")
+        "  integer i\n"
+        "  a({0}) = w({0})*a({0})+x*b({0})+y({0})*c({0})+d({0})*z\n".format(
+            index_str))
     active_variables = ["a", "b", "c", "d"]
     ad_fortran = (
         "  real, dimension(10) :: a\n  real, dimension(10) :: b\n"
         "  real, dimension(10) :: c\n  real, dimension(10) :: d\n"
         "  real, dimension(10) :: w\n  real :: x\n"
-        "  real, dimension(10) :: y\n  real :: z\n\n"
-        "  b(1) = b(1) + x * a(1)\n"
-        "  c(1) = c(1) + y(1) * a(1)\n"
-        "  d(1) = d(1) + a(1) * z\n"
-        "  a(1) = w(1) * a(1)\n\n")
+        "  real, dimension(10) :: y\n  real :: z\n"
+        "  integer :: i\n\n"
+        "  b({0}) = b({0}) + x * a({0})\n"
+        "  c({0}) = c({0}) + y({0}) * a({0})\n"
+        "  d({0}) = d({0}) + a({0}) * z\n"
+        "  a({0}) = w({0}) * a({0})\n\n".format(index_str))
     check_adjoint(tl_fortran, active_variables, ad_fortran, tmpdir)
 
 
-def test_multi_increment(tmpdir):
+def test_multi_increment(tmpdir, index_str):
     '''Test the code works when the active variable that is written to on
     the lhs (A) is also read on the rhs, but is read more than once in
     different terms, some of which are scaled. The resultant adjoint
-    code has the same form as the tangent-linear code.
+    code has the same form as the tangent-linear code. Test for both
+    individual array accesses and array ranges.
 
     A=A+xA -> A=A+xA.
 
     '''
     tl_fortran = (
-        "  real a(10)\n"
-        "  real x\n"
-        "  a(1) = a(1)+x*a(1)\n")
+        f"  real a(10)\n"
+        f"  real x\n"
+        f"  integer i\n"
+        f"  a({index_str}) = a({index_str})+x*a({index_str})\n")
     active_variables = ["a"]
     ad_fortran = (
-        "  real, dimension(10) :: a\n"
-        "  real :: x\n\n"
-        "  a(1) = a(1) + x * a(1)\n\n")
+        f"  real, dimension(10) :: a\n"
+        f"  real :: x\n"
+        f"  integer :: i\n\n"
+        f"  a({index_str}) = a({index_str}) + x * a({index_str})\n\n")
     check_adjoint(tl_fortran, active_variables, ad_fortran, tmpdir)
 
 
@@ -500,7 +589,7 @@ def test_multi_valued_sub(tmpdir):
     check_adjoint(tl_fortran, active_variables, ad_fortran, tmpdir)
 
 
-def test_inc_sub(tmpdir):
+def test_inc_sub(tmpdir, index_str):
     '''Test the transformation works when the active variable that is
     written to on the lhs is also read (and multiplied by -1) on the
     rhs. There are no other active variables on the rhs.
@@ -509,25 +598,26 @@ def test_inc_sub(tmpdir):
 
     '''
     tl_fortran = (
-        "  real a(10)\n"
-        "  integer :: i\n"
-        "  a(i) = -a(i)\n")
+        f"  real a(10)\n"
+        f"  integer :: i\n"
+        f"  a({index_str}) = -a({index_str})\n")
     active_variables = ["a"]
     ad_fortran = (
-        "  real, dimension(10) :: a\n"
-        "  integer :: i\n\n"
-        "  a(i) = -a(i)\n\n")
+        f"  real, dimension(10) :: a\n"
+        f"  integer :: i\n\n"
+        f"  a({index_str}) = -a({index_str})\n\n")
     check_adjoint(tl_fortran, active_variables, ad_fortran, tmpdir)
 
 
-def test_multi_inc_sub(tmpdir):
+def test_multi_inc_sub(tmpdir, index_str):
     '''Test the transformation works when the active variable that is
     written to on the lhs (A) is also read multiple times from
     different terms on the rhs. The terms on the rhs combine the
     active variable with an inactive variable (or no inactive
     variable) in different supported combinations (multiplication,
     division, unary minus). The terms are also combined on the rhs
-    with both addition and subtraction.
+    with both addition and subtraction. Test for both individual array
+    accesses and array ranges.
 
     A=-A-xA+B+A/y -> B*=B*+A*; A*=A*(-1.0-x+1.0/y)
 
@@ -536,14 +626,14 @@ def test_multi_inc_sub(tmpdir):
         "  real a(10), b(10)\n"
         "  integer :: i\n"
         "  real :: x,y\n"
-        "  a(i) = -a(i)-x*a(i)+b(i)+a(i)/y\n")
+        "  a({0}) = -a({0})-x*a({0})+b({0})+a({0})/y\n".format(index_str))
     active_variables = ["a", "b"]
     ad_fortran = (
         "  real, dimension(10) :: a\n  real, dimension(10) :: b\n"
         "  integer :: i\n"
         "  real :: x\n  real :: y\n\n"
-        "  b(i) = b(i) + a(i)\n"
-        "  a(i) = -a(i) - x * a(i) + a(i) / y\n\n")
+        "  b({0}) = b({0}) + a({0})\n"
+        "  a({0}) = -a({0}) - x * a({0}) + a({0}) / y\n\n".format(index_str))
     check_adjoint(tl_fortran, active_variables, ad_fortran, tmpdir)
 
 
@@ -725,6 +815,65 @@ def test_unary_binary_minus(tmpdir):
         "  a = 0.0\n\n")
     check_adjoint(tl_fortran, active_variables, ad_fortran, tmpdir)
 
+
+# _array_ranges_match method
+
+def test_array_ranges_match():
+    ''' Test for the helper method that checks whether the array ranges match
+    for a variable referenced on both the LHS and RHS of an assignment. '''
+    array_type = ArrayType(REAL_TYPE, [ArrayType.Extent.ATTRIBUTE])
+    lhs_symbol = DataSymbol("a", array_type)
+    rhs_symbol = DataSymbol("x", REAL_TYPE)
+    single_elem_ref = ArrayReference.create(lhs_symbol,
+                                            [Literal("1", INTEGER_TYPE)])
+    multiply = BinaryOperation.create(
+        BinaryOperation.Operator.MUL, Reference(rhs_symbol), single_elem_ref)
+    lhs = ArrayReference.create(lhs_symbol,
+                                [Range.create(Literal("2", INTEGER_TYPE),
+                                              Literal("5", INTEGER_TYPE))])
+    assign = Assignment.create(lhs, multiply)
+    trans = AssignmentTrans(active_variables=[lhs_symbol])
+    # If the rhs reference is not to the symbol on the LHS then the routine
+    # should do nothing.
+    trans._array_ranges_match(assign, Reference(rhs_symbol))
+    # If the rhs reference *is* to the same symbol as on the LHS but is not
+    # an ArrayReference then an exception should be raised.
+    with pytest.raises(TangentLinearError) as err:
+        trans._array_ranges_match(assign, Reference(lhs_symbol))
+    assert ("TangentLinearError: Assignment is to an array range but found a "
+            "reference to the LHS variable 'a' without array notation on the "
+            "RHS: 'a(2:5) = x * a(1)" in str(err.value))
+    # If the rhs reference does not match the range on the LHS then an
+    # exception should be raised.
+    with pytest.raises(NotImplementedError) as err:
+        trans._array_ranges_match(assign, single_elem_ref)
+    assert ("Different sections of the same active array 'a' are accessed on "
+            "the LHS and RHS of an assignment: 'a(2:5) = x * a(1)\n'. This "
+            "is not supported." in str(err.value))
+    # Repeat but with a range instead of a single array element. (This isn't
+    # valid code but is sufficient for this test.)
+    wrong_range = Range.create(Literal("3", INTEGER_TYPE),
+                               Literal("6", INTEGER_TYPE))
+    new_ref = ArrayReference.create(lhs_symbol, [wrong_range])
+    assign = Assignment.create(lhs.detach(), new_ref)
+    with pytest.raises(NotImplementedError) as err:
+        trans._array_ranges_match(assign, new_ref)
+    assert ("Different sections of the same active array 'a' are accessed on "
+            "the LHS and RHS of an assignment: 'a(2:5) = a(3:6)\n'. This "
+            "is not supported." in str(err.value))
+    # Repeat with a range where only the step size is wrong. (This isn't valid
+    # code but is sufficient for this test.)
+    wrong_range = Range.create(Literal("2", INTEGER_TYPE),
+                               Literal("5", INTEGER_TYPE),
+                               Literal("2", INTEGER_TYPE))
+    new_ref = ArrayReference.create(lhs_symbol, [wrong_range])
+    assign = Assignment.create(lhs.detach(), new_ref)
+    with pytest.raises(NotImplementedError) as err:
+        trans._array_ranges_match(assign, new_ref)
+    assert ("Different sections of the same active array 'a' are accessed on "
+            "the LHS and RHS of an assignment: 'a(2:5) = a(2:5:2)\n'. This "
+            "is not supported." in str(err.value))
+
 # validate() method
 
 
@@ -1106,6 +1255,66 @@ def test_validate_unaryop():
             "linear with respect to the active variable, but found 'SQRT(b)'."
             in str(info.value))
 
+
+def test_validate_mismatched_array_ranges():
+    ''' Check that we reject a tangent-linear assignment if it references
+    different ranges of the same active variable, e.g:
+                      a(2:5) = x*a(1)
+
+    '''
+    array_type = ArrayType(REAL_TYPE, [ArrayType.Extent.ATTRIBUTE])
+    lhs_symbol = DataSymbol("a", array_type)
+    rhs_symbol = DataSymbol("x", REAL_TYPE)
+    multiply = BinaryOperation.create(
+        BinaryOperation.Operator.MUL, Reference(
+            rhs_symbol), ArrayReference.create(lhs_symbol,
+                                               [Literal("1", INTEGER_TYPE)]))
+    lhs = ArrayReference.create(lhs_symbol,
+                                [Range.create(Literal("2", INTEGER_TYPE),
+                                              Literal("5", INTEGER_TYPE))])
+    assign = Assignment.create(lhs, multiply)
+    trans = AssignmentTrans(active_variables=[lhs_symbol])
+    with pytest.raises(NotImplementedError) as err:
+        trans.validate(assign)
+    assert ("Different sections of the same active array 'a' are accessed on "
+            "the LHS and RHS of an assignment: 'a(2:5) = x * a(1)\n'"
+            in str(err.value))
+    # Repeat but without a multiplier: a(2:5) = a(3:6).
+    assign = Assignment.create(lhs.detach(),
+                               ArrayReference.create(
+                                   lhs_symbol, [Range.create(
+                                       Literal("3", INTEGER_TYPE),
+                                       Literal("6", INTEGER_TYPE))]))
+    trans = AssignmentTrans(active_variables=[lhs_symbol])
+    with pytest.raises(NotImplementedError) as err:
+        trans.validate(assign)
+    assert ("Different sections of the same active array 'a' are accessed on "
+            "the LHS and RHS of an assignment: 'a(2:5) = a(3:6)\n'"
+            in str(err.value))
+
+
+def test_validate_missing_array_indices():
+    ''' Check that we reject a tangent-linear assignment if the LHS is an
+    array access but the same symbol is referenced on the RHS without array
+    notation. (The PSyIR permits this since arrays can be arguments to
+    query functions.) '''
+    array_type = ArrayType(REAL_TYPE, [ArrayType.Extent.ATTRIBUTE])
+    lhs_symbol = DataSymbol("a", array_type)
+    rhs_symbol = DataSymbol("x", REAL_TYPE)
+    multiply = BinaryOperation.create(
+        BinaryOperation.Operator.MUL, Reference(rhs_symbol),
+        Reference(lhs_symbol))
+    lhs = ArrayReference.create(lhs_symbol,
+                                [Range.create(Literal("2", INTEGER_TYPE),
+                                              Literal("5", INTEGER_TYPE))])
+    assign = Assignment.create(lhs, multiply)
+    trans = AssignmentTrans(active_variables=[lhs_symbol])
+    with pytest.raises(TangentLinearError) as err:
+        trans.validate(assign)
+    assert ("Assignment is to an array range but found a reference to the LHS "
+            "variable 'a' without array notation on the RHS: "
+            "'a(2:5) = x * a\n'" in str(err.value))
+
 # _split_nodes() method